sparse.conv export onnx(讲解) - centerpoint

选与握

于 2025-05-30 21:24:08 发布

阅读量960

点赞数 18

CC 4.0 BY-SA版权

文章标签： python 开发语言

本文链接：https://blog.csdn.net/xihuanniNI/article/details/148340735

1 @函数装饰器

装饰器允许你在不修改原始代码的情况下，给函数增加额外的功能。装饰器本质上是一个函数，它接受一个函数作为参数并返回一个新的函数。

demo



# @装饰器，python中顾名思义。和装饰器设计概念一样，这里使用@其实是语法糖
def mydecorator(f):  # 装饰器函数，接受一个函数作为参数, 即被装饰的函数会被作为参数传递给装饰器函数
    def wrapper(*args,**kwargs):  # 接受任意数量和类型的参数，执行被装饰的函数，返回了被装饰的函数，再把参数传递给被装饰的函数，最终执行
        print("装饰器开始")
        f(*args,**kwargs)   # 最终执行被装饰的函数
        print("装饰器结束")
    return wrapper
 
@mydecorator
def hello(name):
    print("hello",name)
 
hello("world")  
# 执行情况： 
# 贴装饰器后，执行hello（“world”）时，hello(name)作为参数传入装饰器函数，装饰器函数执行，返回装饰器实际增强函数，
# 然后传入参数 ‘world’ 即wrapper("world")，wrapper函数执行，执行wrapper函数体内的内容，执行了被装饰的函数
 
# # 装饰器可以叠加使用
# @mydecorator
# @mydecorator   # 从上到下依次装饰，然后从内到外依次执行
# def hello(name):
#     print("hello",name)
 

print('----------------------------------------------------')

# 装饰器可以带参数
def mydecorator2(arg):  # arg-装饰器参数 "hello2"
    def decorator(f):  # 接收要被装饰函数的函数
        def wrapper(*args,**kwargs):  # 如果被装饰的函数带参数则需要这样写，接收被装饰函数的参数
            print("装饰器开始2 arg ",arg) 
            f(*args,**kwargs)   # 执行被装饰的函数
            print("装饰器结束2")  
        return wrapper
    return decorator
 
@mydecorator2("hello2")
def hello2(name):
    print("hello2",name)
 
hello2("world2")

装饰器开始
hello world
装饰器结束
----------------------------------------------------
装饰器开始2 arg hello2
hello2 world2
装饰器结束2

demo2 带参数的原函数

如果原函数是带参数的话，那么装饰器需要怎么写呢？我们需要把参数写到log_decorator里面的那个函数实际上wrapper，下面我们可以通过使用*args和**kwargs来接收和传递参数下面是一个示例，演示如何创建一个适用于带参数的函数的装饰器：

import time

def log_decorator(func):
    def wrapper(*args, **kwargs):
        start_time = time.time()
        func(*args, **kwargs)
        end_time = time.time()
        execution_time = end_time - start_time
        print(f"Function {func.__name__} executed in {execution_time} seconds")
    return wrapper

在上述示例中，log_decorator是一个装饰器函数，它接收任意类型和数量的参数，并使用*args和**kwargs来接收和传递参数。在wrapper函数内部，使用func(*args, **kwargs)来调用原始函数，并将参数传递给它。现在，我们可以使用@语法来应用这个装饰器，无论带有参数还是不带参数的函数都可以。

@log_decorator
def say_hello():
    print("Hello, world!")

@log_decorator
def greet(name):
    print(f"Hello, {name}!")

say_hello()  # 执行带装饰器的say_hello函数
greet("Alittle")  # 执行带装饰器的greet函数

@log_decorator分别为say_hello和greet函数应用了装饰器。无论是不带参数的say_hello函数还是带参数的greet函数，装饰器都能正常工作。

demo3 带参数的装饰器

当需要给装饰器传递参数时，可以使用装饰器工厂函数来创建带参数的装饰器。装饰器工厂函数实际上是一个闭包函数，它接收参数并返回一个真正的装饰器函数。比如我们需要为不同的业务逻辑添加不同的日志等级，就需要在装饰器中添加参数了。

import time

def log_decorator_with_params(log_level):
    def log_decorator(func):
        def wrapper(*args, **kwargs):
            start_time = time.time()
            result = func(*args, **kwargs)
            end_time = time.time()
            execution_time = end_time - start_time
            print(f"{log_level}: Function {func.__name__} executed in {execution_time} seconds")
            return result
        return wrapper
    return log_decorator

如上所示，这种写法可以称为闭包的闭包，log_decorator_with_params是一个装饰器工厂函数，它接收一个参数log_level，用于指定日志的前缀。它返回一个真正的装饰器函数log_decorator，该装饰器函数在函数执行前后打印带有指定前缀的日志。

现在，我们可以使用@语法来应用带参数的装饰器。

@log_decorator_with_params(log_level="INFO")
def say_hello():
    print("Hello, world!")

@log_decorator_with_params(log_level="DEBUG")
def greet(name):
    print(f"Hello, {name}!")

say_hello()  # 执行带装饰器的say_hello函数
greet("Alice")  # 执行带装饰器的greet函数

上述示例中，@log_decorator_with_params("INFO")和@log_decorator_with_params("DEBUG")分别为say_hello和greet函数应用了带参数的装饰器。运行代码时，会分别打印日志等级。在定义装饰器函数wrapper时，使用了*args和**kwargs作为参数，这样能够适配任意类型和数量的参数，并将其传递给原始函数。这样可以确保带参数的函数装饰器适用于不同的函数签名。通过使用装饰器工厂函数，我们可以轻松创建带参数的装饰器，提供更大的灵活性，让装饰器可以根据不同的场景和需求来定制其行为。

demo4

def register_node(fn):

    fnnames   = fn.split(".")
    fn_module = eval(".".join(fnnames[:-1]))
    fn_name   = fnnames[-1]
    oldfn = getattr(fn_module, fn_name)
    
    def make_hook(bind_fn):

        ilayer = 0
        def internal_forward(self, *args):
            global enable_trace

            if not enable_trace:
                return oldfn(self, *args)

            global avoid_reuse_container
            nonlocal ilayer

            # Use the enable_trace flag to avoid internal trace calls
            enable_trace = False
            y = oldfn(self, *args)
            bind_fn(self, ilayer, y, *args) 
            enable_trace = True

            avoid_reuse_container.extend(list(args) + [y]) 
            ilayer += 1
            return y

        setattr(fn_module, fn_name, internal_forward)
    return make_hook

@register_node("spconv.conv.SparseConvolution.forward")
def symbolic_sparse_convolution(self, ilayer, y, x):
    register_tensor(y)
    print(f"   --> SparseConvolution{ilayer}[{'subm' if self.subm else 'conv'}] -> Input {get_tensor_id(x)}, Output {get_tensor_id(y)}")

    if self.transposed:
        output_size = spconv.ops.get_deconv_output_size(
            x.features.size(), self.kernel_size, self.stride, self.padding, self.dilation, self.output_padding
        )
    else:
        output_size = spconv.ops.get_conv_output_size(
            x.features.size(), self.kernel_size, self.stride, self.padding, self.dilation
        )
    
    if self.subm:
        output_size[0] = x.features.size(0)
    
    output_size[1] = self.out_channels
    inputs = [
        get_tensor_id(x), 
        append_initializer(self.weight.data.permute(4, 0, 1, 2, 3), f"spconv{ilayer}.weight"),
    ]

    if self.bias is not None:
        inputs.append(append_initializer(self.bias.data, f"spconv{ilayer}.bias"))

    output_bound = 200000
    if hasattr(self, "output_bound"):
        output_bound = self.output_bound

    nodes.append(
        helper.make_node(
            "SparseConvolution", inputs, [get_tensor_id(y)], f"conv{ilayer}", 
            ndim = self.ndim,
            input_spatial_shape = x.spatial_shape,
            output_spatial_shape = y.spatial_shape,
            in_channels = self.in_channels,
            out_channels = self.out_channels,
            kernel_size = self.kernel_size,
            output_bound = output_bound,
            stride = self.stride,
            dilation = self.dilation,
            padding = self.padding,
            transposed = self.transposed,
            inverse = self.inverse,
            output_padding = self.output_padding,
            groups = self.groups,
            subm = self.subm,
            rulebook = self.indice_key,
            activation = getattr(self, "act_type", "None"),
            input_shape  = x.features.shape,
            output_shape = y.features.shape
        )
    )

下面我们详细分析一下在执行 spconv.conv.SparseConvolution.forward 方法时，代码是如何工作的，主要是涉及到装饰器 register_node 和函数 symbolic_sparse_convolution 的使用

我们先来过一下装饰器 register_node 的工作原理

1. 初始化阶段
- 当 export_tool.py 脚本运行时，首先会去执行 register_node(“spconv.conv.SparseConvolution.forward”) 装饰器
- 这个参数接收一个字符串参数，代表要修改的函数的名称
2. register_node 功能
- 在 register_node 内部，它首先会根据传入的字符串找到相应的函数对象
- 然后，它会返回一个名为 make_hook 的函数
3. 创建 mask_hook
- make_hook 是一个内部函数，它的目的是接收一个函数（这里是 symbolic_sparse_convolution）并“钩住”（hook）原始的 SparseConvolution.forward 方法

接着我们再过一下整个执行过程

1. 应用装饰器
- 当 @register_node(“spconv.conv.SparseConvolution.forward”) 应用于 symbolic_sparse_convolution 函数时，实际上是一个语法糖，其作用等同于 symbolic_sparse_convolution = register_node(“spconv.conv.SparseConvolution.forward”)(symbolic_sparse_convolution)
- make_hook 接收 symbolic_sparse_convolution 作为参数，然后在 SparseConvolution.forward 上调用 internal_forward 函数
2. 修改 SparseConvolution.forward
- internal_forward 函数替换了原始的 SparseConvolution.forward 方法，这意味着每次调用 SparseConvolution.forward 时，实际上是调用 internal_forward
3. 执行 internal forward
- 当 SparseConvolution.forward 被调用时，internal_forward 被执行
- 在 inernal_forward 内部，首先会调用原始的 SparseConvolution.forward 方法（通过 oldfn(self *args)），然后会调用 symbolic_sparse_convolution 函数
4. 执行 symbolic_sparse_convolution
- symbolic_sparse_convolution 函数在每次 SparseConvolution.forward 被调用后执行，用于执行追踪逻辑，如创建 onnx 图中的节点

总的来说，当 spconv.conv.SparseConvolution.forward 方法被调用时，由于装饰器 register_node 的作用，实际上是先执行了 internal_forward 函数。这个函数首先执行原始的 SparseConvolution.forward 方法，然后执行 symbolic_sparse_convolution 函数来进行追踪和收集信息。这种机制允许在不更改原始函数代码的前提下，增加额外的功能，这在许多情况下非常有用，特别是在需要追踪或记录函数行为的场景中。

2. 稀疏卷积概述

spconv 的原理，它通过 hash table 将输入输出中不为 0 的点的位置坐标保存下来，并通过 rulebook 记录输入输出和 filter 之间的对应关系，把没有意义的点全部 skip 掉，只保留真正想做计算的点，从而将一个稀疏的卷积计算变成密集的矩阵乘法计算，实现加速

构建 hash table

第一步是来构建 hash tables

构建 Rulebook

第二步是建立 Rulebook，这是稀疏卷积的关键部分。Rulebook 的目的类似于 im2col，它将卷积从数学形式转换为有效的可编程形式。但与 im2col 不同的是，Rulebook 收集了卷积中所有涉及的原子操作，然后将它们关联到相应的 kernel 元素上。

3. SCN导出

关于稀疏卷积的导出有以下几点说明：

在实现上，traveller59 的 SparseConv 的实现比较完善

https://github.com/traveller59/spconv

我们通常以 CenterPoint 的 SCN 导出为案例来讲解 spconv

https://github.com/tianweiy/CenterPoint/blob/master/det3d/models/backbones/scn.py

SPConv问题的解决方案思维导图如下所示：

图3-1 SPConv问题的解决方案

3.1 实现trace

Trace0

由于 SparseConv 的特殊性，输入输出采用特殊的 tensor 表示。因此标准的 onnx 导出已然无法处理这种复杂的情况

此时可以利用 python 最核心的特性，直接替换特定函数的实现，以实现挂钩到自己函数中

这种做法没有局限性，比 register_forward_hook 要求更低，它可以替换任意函数，而 register_forward_hook 不行

示例代码如下：

import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self):
        super().__init__()

        self.conv1 = nn.Conv2d(3, 3, 1, 1)
        self.relu1 = nn.ReLU()
        self.conv2 = nn.Conv2d(3, 1, 1, 1)
    
    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.conv2(x)
        return x
    
def hook_forward(oldfn):
    def myforward(self, x):
        y = oldfn(self, x)
        print(f"{type(self)} -> Input {id(x)}, Output {id(y)}")
        return y
    return myforward

nn.Conv2d.forward = hook_forward(nn.Conv2d.forward)
nn.ReLU.forward   = hook_forward(nn.ReLU.forward)

model = Model()
x = torch.zeros(1, 3, 3, 3)
y = model(x)

运行效果输入如下：

图3-2 export0输出
在上面的示例代码中，我们通过替换特定函数的实现，实现了钩子函数来挂钩到自定义函数中。具体来说，我们定义了一个名为 hook_forward 的函数，用于替换原有的前向传播函数实现。该函数接受一个旧的前向传播函数作为输入，并返回一个新的前向传播函数。

通过这种方式，我们可以在模型的前向传播过程中插入自定义的操作，例如打印张量的地址、收集统计信息等。它能够保持原有的特性，并嫁接到forward中，实现自定义行为

Trace1

前面我们为了储存旧的 function，采用了闭包特性。这里我们进一步简化 hook 过程，采用字符串解析，加上装饰器。

为了避免 pytorch 对 tensor 进行复用，导致存在 id 相同的 tensor。我们使用了 clone，但是 clone 并不能总是保证唯一，所有这里其实留了一个问题！

示例代码如下：

import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self):
        super().__init__()

        self.conv1 = nn.Conv2d(3, 3, 1, 1)
        self.relu1 = nn.ReLU()
        self.conv2 = nn.Conv2d(3, 1, 1, 1)
    
    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.conv2(x)
        return x
    
def hook_forward(fn):
    
    fnnames   = fn.split(".")
    fn_module = eval(".".join(fnnames[:-1]))
    fn_name   = fnnames[-1]
    oldfn     = getattr(fn_module, fn_name)

    def make_hook(bind_fn):
        def myforward(self, x):
            y = oldfn(self, x).clone()
            bind_fn(self, x, y)
            return y

        setattr(fn_module, fn_name, myforward)
    return make_hook

@hook_forward("torch.nn.Conv2d.forward")
def symbolic_conv2d(self, x, y):
    print(f"{type(self)} -> Input {id(x)}, Output {id(y)}")

@hook_forward("torch.nn.ReLU.forward")
def symbolic_relu(self, x, y):
    print(f"{type(self)} -> Input {id(x)}, Output {id(y)}")
    
model = Model()
x = torch.zeros(1, 3, 3, 3)
y = model(x)

图3-3 export1输出
在上述示例代码中，我们定义了一个 hook_forward 的装饰器函数，用于替换指定函数的前向传播实现。hook_forward 函数接受一个字符串参数 fn，该字符串表示需要替换的函数的路径。通过字符串解析和 eval 函数，获取到需要替换的旧函数对象 oldfn

hook_forward 函数内部定义了一个嵌套函数 make_hook，用于创建实际的钩子函数。make_hook 函数接受一个绑定函数 bind_fn 作为参数，用于在钩子函数中执行额外的操作。其中，bind_fn 函数的形式为 bind_fn(self, x, y)，用于处理输入和输出张量。使用 setattr 函数将新的前向传播函数 myforward 替换旧函数 oldfn，以实现钩子函数的绑定。

我们再来回顾下之前 AutoCV 课程中讲到的关于装饰器的相关知识

Python 装饰器本质上是一个函数，它接受一个函数对应作为参数，并返回一个修改后的函数对象。装饰器通常使用 @decorator 的语法糖来使用，它可以将装饰器应用于函数或类的定义之前，从而实现对其功能的增强或修改。

装饰器的语法结构如下：

@表达式
def 被修饰的函数

其中，表达式需要返回一个函数对象，这个函数对象就是用来修饰函数的。

@hook_forward("torch.nn.Conv2d.forward") 等价于如下代码：

hook_forward("torch.nn.Conv2d.forward")(symbolic_conv2d)

hook_forward("torch.nn.Conv2d.forward") 返回了一个内部函数 make_hook，然后将 symbolic_conv2d 作为参数传递给 make_hook 函数。

由于装饰器本质上是一个函数，它在定义被修饰函数之后立即执行。因此，当定义装饰器 @hook_forward("torch.nn.Conv2d.forward") 时，装饰器函数 hook_forward 会被调用，并且 setattr 语句会在函数内部执行。此时，make_hook(symbolic_conv2d) 将创建一个新的前向传播函数，并将其替换为 nn.Conv2d 的 forward 函数，从而实现了将钩子函数绑定到前向传播函数上的目的。

Trace2

在 Trace1 中我们提到需要避免 pytorch 对 tensor 进行复用，下面我们就来解决它

首先思考下 pytorch 在何时会复用 tensor？

答案是在没有任何引用的 tensor 会被回收并复用

那么解决方案就是引用它，不释放，然后把 id 重新编号为 tensor 数

示例代码如下：

import torch
import torch.nn as nn

class Model(nn.Module):
    def __init__(self):
        super().__init__()

        self.conv1 = nn.Conv2d(3, 3, 1, 1)
        self.relu1 = nn.ReLU()
        self.conv2 = nn.Conv2d(3, 1, 1, 1)
    
    def forward(self, x):
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.conv2(x)
        return x
    
def hook_forward(fn):

    fnnames   = fn.split(".") 
    fn_module = eval(".".join(fnnames[:-1]))
    fn_name   = fnnames[-1]
    oldfn = getattr(fn_module, fn_name)

    def make_hook(bind_fn):
        def myforward(self, x):
            global all_tensors
            y = oldfn(self, x)
            bind_fn(self, x, y)
            all_tensors.extend([x, y])  # 避免torch对tensor进行复用
            return y
        
        setattr(fn_module, fn_name, myforward)
    return make_hook

@hook_forward("torch.nn.Conv2d.forward")
def symbolic_conv2d(self, x, y):
    print(f"{type(self)} -> Input {get_obj_idd(x)}, Output {get_obj_idd(y)}")

@hook_forward("torch.nn.ReLU.forward")
def symbolic_relu(self, x, y):
    print(f"{type(self)} -> Input {get_obj_idd(x)}, Output {get_obj_idd(y)}")

def get_obj_idd(obj):
    global objmap

    idd = id(obj)
    if idd not in objmap:
        objmap[idd] = len(objmap)
    return objmap[idd]

# 避免torch对内存复用导致id相同
all_tensors = []

# 为每个新的tensor编号
objmap = {}

model = Model()
x = torch.zeros(1, 3, 3, 3)
y = model(x)

图3-4 export2输出
上述示例代码通过引用 Tensor 并重新编号其内存地址的方式，避免了 PyTorch 对 Tensor 进行复用。我们定义了一个 get_obj_idd 函数，用于为每个 Tensor 对象分配一个唯一的编号，避免 PyTorch 对内存的复用导致 id 相同。其中空列表 all_tensors 用于存储所有的 Tensor 对象，空字典 objmap 用于映射 Tensor 对象的内存地址到其编号。

到此为止 Trace 的核心功能得到实现，剩下的是把 trace 的 graph 交给 onnx

3.2 导出onnx

由于 trace 是我们自己实现的，因此 onnx 的创建工作也需要我们自己来动手

关于 onnx 的操作可参考 https://shouxieai.com/solution/trt/basic-1.4-onnx-editor

在开始之前我们还是需要了解下 onnx 文件的组成，方便后续操作，onnx 文件组成如下图所示：

图3-5 onnx文件组成
model：表示整个 onnx 模型，包括图结构和解析器版本、opset 版本、导出程序类型
opset 版本即 operator 版本号即 pytorch 的 op(操作算子) 版本
model.graph：表示图结构，通常是 Netron 中看到的结构
model.graph.node：表示图结构中所有节点如 conv、bn、relu 等
model.graph.initializer：权重数据大都存储在这里
model.graph.input：模型的输入，它指定了输入的名称、数据类型和形状
model.graph.output：模型的输出，它指定了输出的名称、数据类型和形状
因此我们要创建一个 onnx 需要定义模型的输入和输出即 model.graph.input 和 model.graph.output，其余的就是一些节点的定义，对应的权重使用 model.graph.initializer 进行初始化即可

下面是利用 onnx.helper 创建一个节点的示例：

helper.make_node(
name="Conv_0", # 节点名字，注意与op_type的区分
op_type="Conv", # 节点的算子类型，比如'Conv'、'Relu'、'Add'
inputs=["image", "conv.weight", "conv.bias"], # 各个输入的名字，节点的输入包含：输入和算子的权重
outputs=["3"], # 输出的个数
pads=[1, 1, 1, 1], # 其他字符串为该节点的属性，比如'Conv'的stride、pad等
group=1,
dilations=[1,1],
kernel_shape=[3,3],
strides=[1,1]
)

自己实现的 trace 的 graph 导出为 onnx 的示例代码如下：

import torch
import torch.nn as nn
import onnx
import onnx.helper as helper
import numpy as np

# reference
# https://shouxieai.com/solution/trt/basic-1.4-onnx-editor

class Model(nn.Module):
    def __init__(self):
        super().__init__()

        self.conv1 = nn.Conv2d(3, 3, 1, 1)
        self.relu1 = nn.ReLU()
        self.conv2 = nn.Conv2d(3, 1, 1, 1)
        self.conv_right = nn.Conv2d(3, 3, 1, 1)
    
    def forward(self, x):
        r = self.conv_right(x)
        x = self.conv1(x)
        x = self.relu1(x)
        x = self.conv2(x + r)
        return x

def hook_forward(fn):
    
    fnnames   = fn.split(".")
    fn_module = eval(".".join(fnnames[:-1]))
    fn_name   = fnnames[-1]
    oldfn = getattr(fn_module, fn_name)

    def make_hook(bind_fn):
        
        ilayer = 0
        def myforward(self, x):
            global all_tensors
            nonlocal ilayer
            y = oldfn(self, x)

            bind_fn(self, ilayer, x, y)
            all_tensors.extend([x, y])
            ilayer += 1
            return y
    
        setattr(fn_module, fn_name, myforward)
    return make_hook

@hook_forward("torch.nn.Conv2d.forward")
def symbolic_conv2d(self, ilayer, x, y):
    print(f"{type(self)} -> Input {get_obj_idd(x)}, Output {get_obj_idd(y)}")

    inputs = [
        get_obj_idd(x),
        append_initializer(self.weight.data, f"conv{ilayer}.weight"),
        append_initializer(self.bias.data, f"conv{ilayer}.bias")
    ]

    nodes.append(
        helper.make_node(
            "Conv", inputs, [get_obj_idd(y)], f"conv{ilayer}",
            kernel_shape=self.kernel_size, group=self.groups, 
            pads=[0, 0] + list(self.padding), dilations=self.dilation, strides=self.stride
        )
    )

@hook_forward("torch.nn.ReLU.forward")
def symbolic_relu(self, ilayer, x, y):
    print(f"{type(self)} -> Input {get_obj_idd(x)}, Output {get_obj_idd(y)}")

    nodes.append(
        helper.make_node(
            "Relu", [get_obj_idd(x)], [get_obj_idd(y)], f"relu{ilayer}"
        )
    )

@hook_forward("torch.Tensor.__add__")
def symbolic_add(a, ilayer, b, y):
    print(f"Add -> Input {get_obj_idd(a)} + {get_obj_idd(b)}, Output {get_obj_idd(y)}")

    nodes.append(
        helper.make_node(
            "Add", [get_obj_idd(a), get_obj_idd(b)], [get_obj_idd(y)], f"add{ilayer}"
        )
    )

def append_initializer(value, name):
    initializers.append(
        helper.make_tensor(
            name=name,
            data_type=helper.TensorProto.DataType.FLOAT,
            dims=list(value.shape),
            vals=value.data.numpy().astype(np.float32).tobytes(),
            raw=True
        )
    )
    return name

def get_obj_idd(obj):
    global objmap
    
    idd = id(obj)
    if idd not in objmap:
        objmap[idd] = str(len(objmap))
    return objmap[idd]

all_tensors = []
objmap = {}
nodes = []
initializers = []

torch.manual_seed(31)
x = torch.full((1, 3, 3, 3), 0.55)
model = Model().eval()
y = model(x)

inputs = [
    helper.make_value_info(
        name="0",
        type_proto=helper.make_tensor_type_proto(
            elem_type=helper.TensorProto.DataType.FLOAT,
            shape=["batch", x.size(1), x.size(2), x.size(3)]
        )
    )
]

outputs = [
    helper.make_value_info(
        name="5",
        type_proto=helper.make_tensor_type_proto(
            elem_type=helper.TensorProto.DataType.FLOAT,
            shape=["batch", y.size(1), y.size(2), y.size(3)]
        )
    )
]

graph = helper.make_graph(
    name="mymodel",
    inputs=inputs,
    outputs=outputs,
    nodes=nodes,
    initializer=initializers
)

# 如果名字不是ai.onnx，netron解析就不是太一样了
opset = [
    helper.make_operatorsetid("ai.onnx", 11)
]

# producer主要是保持与pytorch一致
model = helper.make_model(graph, opset_imports=opset, producer_name="pytorch", producer_version="1.12")
onnx.save_model(model, "./custom.onnx")

print(y)

在上述示例代码中，通过 helper.make_node 函数创建节点，并将其添加到 nodes 列表中。append_initializer 函数用于将权重参数添加到 initializers 列表中，以便导出到 onnx 模型中。然后使用 helper.make_graph 创建计算图，传入输入、输出、节点和初始化器信息。接下来利用 helper.make_model 函数创建模型，最后使用 onnx.save_model 将模型保存为 onnx 格式的文件。

运行效果如下：

图3-6 export3输出

导出的 onnx 模型如下：

图3-7 导出的onnx

可以看到导出的 onnx 模型符合我们的预期，由于是导出的 onnx，因此我们可以利用 onnxruntime 对模型进行推理验证，验证的示例代码如下：

import onnxruntime
import numpy as np

session = onnxruntime.InferenceSession("custom.onnx", providers=["CPUExecutionProvider"])

x = np.full((1, 3, 3, 3), 0.55, dtype=np.float32)
y = session.run(["5"], {"0": x})[0]

print(y)

图3-8 onnx模型输出
可以看到 onnxruntime 推理验证的结果与之前的一样，可知整个过程应该没什么问题

3.3 CenterPoint SCN导出

主要代码在 https://github.com/tianweiy/CenterPoint/blob/master/det3d/models/backbones/scn.py 由于他的 forward 函数不规范，我们需要改造他，使得输入输出都是单纯的 sparseConvTensor

原始的 forward 函数如下：

def forward(self, voxel_features, coors, batch_size, input_shape):
    
    # input: # [41, 1600, 1408]
    sparse_shape = np.array(input_shape[::-1]) + [1, 0, 0]
    
    coors = coors.int()
    ret = spconv.SparseConvTensor(voxel_features, coors, sparse_shape, batch_size)
    
    x = self.conv_input(ret)
    
    x_conv1 = self.conv1(x)
    x_conv2 = self.conv2(x_conv1)
    x_conv3 = self.conv3(x_conv2)
    x_conv4 = self.conv4(x_conv3)
    
    ret = self.extra_conv(x_conv4)
    
    ret = ret.dense()
    
    N, C, D, H, W = ret.shape
    ret = ret.view(N, C * D, H, W)
    
    multi_scale_voxel_features = {
        'conv1': x_conv1,
        'conv2': x_conv2,
        'conv3': x_conv3,
        'conv4': x_conv4
    }
    
    return ret, multi_scale_voxel_features

修改后的 forward 函数如下：

def forward(self, x : spconv.SparseConvTensor):
    
    # input: # [41, 1600, 1408]
    # sparse_shape = np.array(input_shape[::-1]) + [1, 0, 0]
    
    # coors = coors.int()
    # ret = spconv.SparseConvTensor(voxel_features, coors, sparse_shape, batch_size)
    
    x = self.conv_input(x)
    
    x_conv1 = self.conv1(x)
    x_conv2 = self.conv2(x_conv1)
    x_conv3 = self.conv3(x_conv2)
    x_conv4 = self.conv4(x_conv3)
    
    ret = self.extra_conv(x_conv4)
    
    # 后续的dense只是一个scatter操作，比较简单。目前先分离他们
    
    return ret

通过加载点云特征和坐标数据，可以写一个 SCN 的简单推理案例。好进一步展开，示例代码如下：

from det3d.models.backbones.scn import SpMiddleResNetFHD
from spconv.pytorch import SparseConvTensor
import torch
import pickle

with open("test_spconv.pkl", "rb") as f:
    (voxels, coors, spatial_shape) = pickle.load(f)

print(voxels.shape, voxels.dtype, coors.shape, coors.dtype, spatial_shape)
model = SpMiddleResNetFHD(voxels.shape[1]).cuda().eval().half()

voxels = torch.from_numpy(voxels).cuda().half()
coors  = torch.from_numpy(coors).cuda()

x = SparseConvTensor(voxels, coors, spatial_shape, 1)

with torch.no_grad():
    y = model(x)

    print(y.features.shape, y.indices.shape, y.spatial_shape)

注意这里用了 half，是因为 spconv 支持 fp16 推理，性能比较好。coors 就是 indices，包括 [batch,x,y,z] 四个维度，它是 int 类型的

我们现在可以将之前的 trace 对接到 SCN 上了，有以下几点值得我们注意：

在 myforward 中我们增加了 enable_trace 标记保护，避免在调用 oldfn 时，里面再次触发 hook

def hook_forward(fn):
    ...
    
    def make_hook(bind_fn):
        
        ilayer = 0
        def myforward(self, x):
            global all_tensors, enable_trace
            nonlocal ilayer
            
            # 标记保护
            if not enable_trace: return
            
            enbale_trace = Face
            y = oldfn(self, x)
            
            bind_fn(self, ilayer, x, y)
            enable_trace = True
            all_tensors.extend([x, y])
            ilayer += 1
            return y
        
        setattr(fn_module, fn_name, myforward)
    return make_hook

hook 住 sparseConvolution 的 forward 函数，并添加对应的 node，和之前的 Conv、ReLU 一样。我们这里把 SparseConvTensor 作为一个输入值即可，不用区分 features 和 indices，因为不用走 tensorRT

@hook_forward("spconv.pytorch.conv.SparseConvolution.forward")
def symbolic_conv2d(self, ilayer, x, y):
    print(f"{type(self)} -> Input {get_obj_idd(x)}, Output {get_obj_idd(y)}")
    
    inputs = [
        get_obj_idd(x),
        append_initializer(self.weight.data, f"conv{ilayer}.weight")
    ]
    
    if self.bias is not None:
        inputs.append(append_initializer(self.bias.data, f"conv{ilayer}.bias"))
        
    nodes.append(
    	helper.make_node(
        	"SparseConvolution", inputs, [get_obj_idd(y)], f"spconv{ilayer}",
            kernel_shape=self.kernel_size, group=self.groups, pads=[0, 0] + list(self.padding),
            dilations=self.dilation, strides=self.stride
        )
    )

在 ReLU 节点中需要避免 inplace 操作，因为 inplace 会使得 input、output 的 id 一样

for name, m in model.named_modules():
    if isinstance(m, nn.ReLU):
        m.inplace = False

这里面存在很多的 bn、relu 可以进行融合，这是 spconv 提供的一些支持，其中 bn 可以与 spconv 的 weight、bias 进行 fusion，activation(relu) 则可以与 spconv 的实现进行融合，spconv 可以通过增加 act_type 标记告诉它 activation 是什么。融合后没有 bn 和 relu，更加简单

def fuse_conv_bn(self, conv_out, conv, bn):
    
    # conv ->
    # y = x * w + b

    # bn
    # t = (x - mean) / std
    # t = x * 1/var + (-mean / var)
    # y = t * gamma + beta
    # y = (x * w + b) * 1/var * gamma + (-mean/var) * gamma + beta
    # y = x * w * 1/var * gamma + (-mean/var) * gamma + beta + conv.b * 1/var * gamma
    
    # conv -> bn
    # y1 = x * conv.w + conv.b
    # y2 = (y1 - mean) / var
    # y3 = y2 * gamma + beta
    # output = ((x * conv.w + conv.b) - mean) / var * gamma + beta
    # output = (x * conv.w + conv.b - mean) / var * gamma + beta
    #        = x * conv.w / var * gamma + conv.b / var * gamma - mean / var * gamma + beta
    # weight = x * conv.w / var * gamma
    # bias   = conv.b / var * gamma - mean / var * gamma + beta

    std = torch.sqrt(bn.running_var.data) + bn.eps
    conv_out.weight.data[:] = conv.weight.data / std.view(1, -1, 1, 1) * bn.weight.data.view(1, -1, 1, 1)
    conv_out.bias.data[:]   = conv.bias.data / std * bn.weight.data + bn.bias.data + (-bn.running_mean.data / std) * bn.weight.data

SCN 导出 onnx 的示例代码如下：

from det3d.models.backbones.scn import SpMiddleResNetFHD
from spconv.pytorch import SparseConvTensor
import sponnx
import torch
import pickle
import torch.nn as nn
import onnx
import onnx.helper as helper
import numpy as np
import spconv

torch.manual_seed(31)
with open("test_spconv.pkl", "rb") as f:
    (voxels, coors, spatial_shape) = pickle.load(f)

print(spatial_shape)
model = SpMiddleResNetFHD(voxels.shape[1]).cuda().eval().half()


def hook_forward(fn):

    fnnames   = fn.split(".")
    fn_module = eval(".".join(fnnames[:-1]))
    fn_name   = fnnames[-1]
    oldfn = getattr(fn_module, fn_name)
    
    def make_hook(bind_fn):

        ilayer = 0
        def myforward(self, x):
            global all_tensors, enable_trace
            nonlocal ilayer

            if not enable_trace: return

            enable_trace = False
            y = oldfn(self, x)

            bind_fn(self, ilayer, x, y)
            enable_trace = True
            all_tensors.extend([x, y])   # 避免torch对tensor进行复用
            ilayer += 1
            return y

        setattr(fn_module, fn_name, myforward)
    return make_hook

@hook_forward("spconv.pytorch.conv.SparseConvolution.forward")
def symbolic_conv2d(self, ilayer, x, y):
    print(f"{type(self)} -> Input {get_obj_idd(x)}, Output {get_obj_idd(y)}")

    inputs = [
        get_obj_idd(x),
        append_initializer(self.weight.data, f"conv{ilayer}.weight")
    ]

    if self.bias is not None:
        inputs.append(append_initializer(self.bias.data, f"conv{ilayer}.bias"))

    nodes.append(
        helper.make_node(
            "SparseConvolution", inputs, [get_obj_idd(y)], f"spconv{ilayer}", 
            kernel_shape=self.kernel_size, group=self.groups, pads=[0, 0] + list(self.padding), dilations=self.dilation, strides=self.stride
        )
    )

@hook_forward("torch.nn.ReLU.forward")
def symbolic_relu(self, ilayer, x, y):
    print(f"{type(self)} -> Input {get_obj_idd(x)}, Output {get_obj_idd(y)}")

    nodes.append(
        helper.make_node(
            "Relu", [get_obj_idd(x)], [get_obj_idd(y)], f"relu{ilayer}"
        )
    )

@hook_forward("torch.Tensor.__add__")
def symbolic_add(self, ilayer, x, y):
    print(f"Add -> Input {get_obj_idd(self)} + {get_obj_idd(x)}, Output {get_obj_idd(y)}")

    nodes.append(
        helper.make_node(
            "Add", [get_obj_idd(self), get_obj_idd(x)], [get_obj_idd(y)], f"add{ilayer}"
        )
    )

@hook_forward("torch.nn.BatchNorm1d.forward")
def node_batchnorm1d(self, ilayer, x, y):
    print(f"{type(self)} -> Input {get_obj_idd(x)}, Output {get_obj_idd(y)}")

    nodes.append(
        helper.make_node(
            "BatchNormalization", 
            [  
                get_obj_idd(x), 
                append_initializer(self.weight, f"bn{ilayer}.weight"),
                append_initializer(self.bias, f"bn{ilayer}.bias"),
                append_initializer(self.running_mean, f"bn{ilayer}.running_mean"),
                append_initializer(self.running_var, f"bn{ilayer}.running_var"),
            ],
            [get_obj_idd(y)],
            epsilon=self.eps,
            momentum=self.momentum,
            name=f"batch_norm{ilayer}"
        )
    )

def append_initializer(value, name):
    initializers.append(
        helper.make_tensor(
            name=name,
            data_type=helper.TensorProto.DataType.FLOAT,
            dims=list(value.shape),
            vals=value.data.cpu().numpy().astype(np.float32).tobytes(),
            raw=True
        )
    )
    return name


def get_obj_idd(obj):
    global objmap

    if isinstance(obj, SparseConvTensor):
        obj = obj.features

    idd = id(obj)
    if idd not in objmap:
        objmap[idd] = str(len(objmap))
    return objmap[idd]

enable_trace = True
all_tensors = []
objmap = {}
nodes = []
initializers = []

voxels = torch.from_numpy(voxels).cuda().half()
coors  = torch.from_numpy(coors).cuda()

x = SparseConvTensor(voxels, coors, spatial_shape, 1)

for name, m in model.named_modules():
    if isinstance(m, nn.ReLU):
        m.inplace = False

with torch.no_grad():
    y = model(x)

inputs = [
    helper.make_value_info(
        name="0",
        type_proto=helper.make_tensor_type_proto(
            elem_type=helper.TensorProto.DataType.FLOAT,
            shape=["n", x.features.size(1)]
        )
    )
]

outputs = [
    helper.make_value_info(
        name=nodes[-1].output[0],
        type_proto=helper.make_tensor_type_proto(
            elem_type=helper.TensorProto.DataType.FLOAT,
            shape=["n", y.features.size(1)]
        )
    )
]

graph = helper.make_graph(
    name="mymodel",
    inputs=inputs,
    outputs=outputs,
    nodes=nodes,
    initializer=initializers
)

# 如果名字不是ai.onnx，netron解析就不是太一样了
opset = [
    helper.make_operatorsetid("ai.onnx", 11)
]

# producer主要是保持和pytorch一致
model = helper.make_model(graph, opset_imports=opset, producer_name="pytorch", producer_version="1.9")
onnx.save_model(model, "scn.onnx")

除了上面提到的几点外，其他的和之前说的没有什么差别。

值得注意的是，导出的 onnx 并不能被 tensorRT 使用，因为它的输入和输出是复合类型，而 tensorRT 不支持这种复合类型，所以只能靠自己构建和推理执行图

导出的 onnx 模型如下图所示：

图3-9 SCN的onnx导出

3.4 执行图的构建

前面我们解决了模型的 onnx 导出问题，现在来探讨模型推理问题，由于 tensorRT 这条路走不通，因此只能自行解析 onnx 单独进行推理。

onnx 是静态的计算图，我们需要构建一个执行图(如图3-10所示)，执行图相比静态图最大的区别就是，每个节点都是具有真实值的，每个节点都可以执行具体计算，它是计算图起作用模式下的样子。

图3-10 执行图示例

我们学习 tensorRT 构建执行图的 API，我们可以做如下设计来表示执行图

a = engine.add_input()	# 返回tensor
b = engine.add_input()	# 返回tensor
conv1 = engine.add_conv(a)	# 返回conv layer
conv2 = engine.add_conv(b)	# 返回conv alyer
add = engine.add_add(conv1.output, conv2.output) # 返回 add layer
engine.mark_output(add.output)	# 标记最终需要保留的输出

当我们构建好执行图后

为 a、b 赋值
向 e 索要最新值，即 e.update()
e 的最新值需要通过 e.parent.compute() 得到
add 的 compute 实现为 c 和 d 必须最新值才能计算，因此 c.update()，d.update()
由于 c.update() 需要 c.parent.compute()，因此 Conv1 需要执行 compute，此时 a 已经是最新值，计算后结果给到 c
由于 d.update() 需要 d.parent.compute()，因此 Conv2 需要执行 compute，此时 b 已经是最新值，计算后结果给到 d
然后执行 add 的加法操作，e=c+d
此时，拿到的 e 已经是最新值。作为执行图结果返回

下面是 tensorRT 下的执行图构建方式，我们可以拿来参考

// 构建一个模型
/*
    Network definition:

    image
      |
    linear (fully connected) input = 3, output = 2, bias = True
      |
    sigmoid
      |
    prob
*/

// -------------------------2. 输入，模型结构和输出的基本信息-----------------------
const int num_input = 3;    // in_channel
const int num_output = 2;   // out_channel
float layer1_weight_values[] = {1.0, 2.0, 0.5, 0.1, 0.2, 0.5};  // 前3个给w1的rgb，后3个给w2的rgb
float layer1_bias_values[]   = {0.3, 0.8}

// 输入指定数据的名称、数据类型和完整维度，将输入层添加到网络
nvinfer1::ITensor* input = network->addInput("image", nvinfer1::DataType::kFLOAT, nvinfer1::Dims4(1, num_input, 1, 1));
nvinfer1::Weights layer1_weight = make_weights(layer1_weight_values, 6);
nvinfer1::Weights layer1_bias   = make_weights(layer1_bias_values, 2);
// 添加全连接层
auto layer1 = network->addFullyConnected(*input, num_output, layer1_weight, layer1_bias);   // 注意对input进行了解引用
// 添加激活层
auto prob = network->addActivation(*layer1->getOuput(0), nvinfer1::ActivationType::kSIGMOID);

// 将我们需要的prob标注为输出
network->markOutput(*prob->getOuput(0));

spconv 推理的演示代码如下

import numpy as np

class Tensor:
    def __init__(self, name, parent=None):
        self.name = name
        self.value = 0
        self.parent = parent
    
    def update(self):
        if self.parent is not None:
            self.parent.update()

class Node:
    def __init__(self, name, op_type):
        self.name = name
        self.op_type = op_type
        self.is_computed = False
    
    def update(self):
        if not self.is_computed:
            self.is_computed = True

            for x in self.input:
                x.update()

            self.forward()

class SparseConvolution(Node):
    def __init__(self, name, x, attributes):
        super().__init__(name, "SparseConvolution")
        self.attributes = attributes
        self.input = [x]
        self.output = Tensor(f"{name}.output", self)

    def forward(self):
        self.output.value = self.input[0].value * 0.5
        print(f"Do {self.op_type} x[{self.input[0].value}] * 0.5, output = {self.output.value}")

class BatchNormalization(Node):
    def __init__(self, name, x, attributes):
        super().__init__(name, "BatchNormalization")
        self.input = [x]
        self.attributes = attributes
        self.output = Tensor(f"{name}.output", self)

    def forward(self):
        self.output.value = self.input[0].value + 0.1
        print(f"Do {self.op_type} x[{self.input[0].value}] + 0.1, output = {self.output.value}")


class Engine:
    def __init__(self):
        self.inputs  = []
        self.outputs = []
        self.nodes = []

    def add_input(self, name):
        x = Tensor(name)
        self.inputs.append(x)
        return x

    def mark_output(self, x):
        self.outputs.append(x)
        return x

    def add_spconv(self, name, x, attributes=None):
        x = SparseConvolution(name, x, attributes)
        self.nodes.append(x)
        return x

    def add_bn(self, name, x, attributes=None):
        x = BatchNormalization(name, x, attributes)
        self.nodes.append(x)
        return x

    def forward(self, x):

        for n in self.nodes:
            n.is_computed = False

        self.inputs[0].value = x
        self.outputs[0].update()
        return self.outputs[0].value


engine = Engine()
x = engine.add_input("input")
spconv0 = engine.add_spconv("spconv0", x)
bn1     = engine.add_bn("bn1", spconv0.output)
spconv1 = engine.add_spconv("spconv1", bn1.output)
bn1     = engine.add_bn("bn1", spconv1.output)
pred    = engine.add_spconv("pred", bn1.output)
engine.mark_output(pred.output)

print(engine.forward(1))

运行效果如下：

图3-11 infer0输出
上述代码演示了一个简化的 ONNX 执行图构建过程。其中，Engine 类用于构建和运行计算图，通过添加输入张量和各种操作节点来定义计算图的结构。每个节点通过重写 forward() 方法来定义其具体的计算操作，并通过调用 update() 方法逐层更新计算图中的张量值。通过设置输入张量的值并调用 forward() 方法，可以实现计算图的前向传播，得到输出结果。

3.5 onnx解析并创建执行图

由于 onnx 格式储存的计算图是一个很友好的方式，因此配合好 onnx 的格式与我们写的 add 系列的 api，我们就可以直接实现一个 onnx 解析器，直接解析并放到 engine 里面。

演示的示例代码如下

import numpy as np
import onnx

class Tensor:
    def __init__(self, name, parent=None):
        self.name = name
        self.value = 0
        self.parent = parent
    
    def update(self):
        if self.parent is not None:
            self.parent.update()

class Node:
    def __init__(self, name, op_type):
        self.name = name
        self.op_type = op_type
        self.is_computed = False
    
    def update(self):
        if not self.is_computed:
            self.is_computed = True

            for x in self.input:
                x.update()

            self.forward()

class SparseConvolution(Node):
    def __init__(self, name, x, attributes):
        super().__init__(name, "SparseConvolution")
        self.attributes = attributes
        self.input = [x]
        self.output = Tensor(f"{name}.output", self)

    def forward(self):
        self.output.value = self.input[0].value * 0.5
        print(f"Do {self.op_type} x[{self.input[0].value}] * 0.5, output = {self.output.value}")

class ReLU(Node):
    def __init__(self, name, x):
        super().__init__(name, "ReLU")
        self.input = [x]
        self.output = Tensor(f"{name}.output", self)

    def forward(self):
        self.output.value = max(0, self.input[0].value)
        print(f"Do {self.op_type} max(0, x[{self.input[0].value}]), output = {self.output.value}")


class Add(Node):
    def __init__(self, name, a, b):
        super().__init__(name, "Add")
        self.input = [a, b]
        self.output = Tensor(f"{name}.output", self)

    def forward(self):
        self.output.value = self.input[0].value + self.input[1].value
        print(f"Do {self.op_type} a[{self.input[0].value}] + b[{self.input[1].value}], output = {self.output.value}")


class BatchNormalization(Node):
    def __init__(self, name, x, attributes):
        super().__init__(name, "BatchNormalization")
        self.input = [x]
        self.attributes = attributes
        self.output = Tensor(f"{name}.output", self)

    def forward(self):
        self.output.value = self.input[0].value + 0.1
        print(f"Do {self.op_type} x[{self.input[0].value}] + 0.1, output = {self.output.value}")


class Engine:
    def __init__(self):
        self.inputs  = []
        self.outputs = []
        self.nodes = []

    def add_input(self, name):
        x = Tensor(name)
        self.inputs.append(x)
        return x

    def mark_output(self, x):
        self.outputs.append(x)
        return x

    def add_spconv(self, name, x, attributes=None):
        x = SparseConvolution(name, x, attributes)
        self.nodes.append(x)
        return x

    def add_relu(self, name, x):
        x = ReLU(name, x)
        self.nodes.append(x)
        return x
    
    def add_add(self, name, a, b):
        x = Add(name, a, b)
        self.nodes.append(x)
        return x

    def add_bn(self, name, x, attributes=None):
        x = BatchNormalization(name, x, attributes)
        self.nodes.append(x)
        return x

    def forward(self, x):

        for n in self.nodes:
            n.is_computed = False

        self.inputs[0].value = x
        self.outputs[0].update()
        return self.outputs[0].value


def load_engine(onnxfile):

    model = onnx.load(onnxfile)
    engine = Engine()
    name_to_tensor = {}

    for i in model.graph.input:
        x = engine.add_input(i.name)
        name_to_tensor[x.name] = x
        
    for n in model.graph.node:
        if n.op_type == "SparseConvolution":
            layer = engine.add_spconv(n.name, name_to_tensor[n.input[0]], n)
            name_to_tensor[n.output[0]] = layer.output
        elif n.op_type == "BatchNormalization":
            layer = engine.add_bn(n.name, name_to_tensor[n.input[0]], n)
            name_to_tensor[n.output[0]] = layer.output
        elif n.op_type == "Relu":
            layer = engine.add_relu(n.name, name_to_tensor[n.input[0]])
            name_to_tensor[n.output[0]] = layer.output
        elif n.op_type == "Add":
            layer = engine.add_add(n.name, name_to_tensor[n.input[0]], name_to_tensor[n.input[1]])
            name_to_tensor[n.output[0]] = layer.output
        else:
            raise RuntimeError(f"Unsupport op_type {n.op_type}")

    for o in model.graph.output:
        engine.mark_output(name_to_tensor[o.name])

    return engine    
    

engine = load_engine("scn.onnx")
print(engine.forward(1.0))

运行效果如下：

图3-12 infer1输出

至此，整个 onnx 解析器已经完成。你应该在 C++ 上复现这个流程就可以实现无 trt 推理了

4 bevfusion 输出

    print("Tracing model inference...")
    print("> Do inference...")
    with torch.no_grad():
        register_tensor(voxels)
        enable_trace = True
        y = model(voxels, coors, batch_size)  # 执行这里，调用符号函数
        enable_trace = False

    print("Tracing done!")

Tracing model inference...
> Do inference...
--> SparseConvolutionQunat0[subm] -> Input 0, Output 1
Backend TkAgg is interactive backend. Turning interactive mode on.
--> SparseConvolutionQunat1[subm] -> Input 1, Output 2
--> SparseConvolutionQunat2[subm] -> Input 2, Output 3
--> QuantAdd0 -> Input 3 + 1, Output 4
--> ReLU0 -> Input 4, Output 5
--> SparseConvolutionQunat3[subm] -> Input 5, Output 6
--> SparseConvolutionQunat4[subm] -> Input 6, Output 7
--> QuantAdd1 -> Input 7 + 5, Output 8
--> ReLU1 -> Input 8, Output 9
--> SparseConvolutionQunat5[conv] -> Input 9, Output 10
--> SparseConvolutionQunat6[subm] -> Input 10, Output 11
--> SparseConvolutionQunat7[subm] -> Input 11, Output 12
--> QuantAdd2 -> Input 12 + 10, Output 13
--> ReLU2 -> Input 13, Output 14
--> SparseConvolutionQunat8[subm] -> Input 14, Output 15
--> SparseConvolutionQunat9[subm] -> Input 15, Output 16
--> QuantAdd3 -> Input 16 + 14, Output 17
--> ReLU3 -> Input 17, Output 18
--> SparseConvolutionQunat10[conv] -> Input 18, Output 19
--> SparseConvolutionQunat11[subm] -> Input 19, Output 20
--> SparseConvolutionQunat12[subm] -> Input 20, Output 21
--> QuantAdd4 -> Input 21 + 19, Output 22
--> ReLU4 -> Input 22, Output 23
--> SparseConvolutionQunat13[subm] -> Input 23, Output 24
--> SparseConvolutionQunat14[subm] -> Input 24, Output 25
--> QuantAdd5 -> Input 25 + 23, Output 26
--> ReLU5 -> Input 26, Output 27
--> SparseConvolutionQunat15[conv] -> Input 27, Output 28
--> SparseConvolutionQunat16[subm] -> Input 28, Output 29
--> SparseConvolutionQunat17[subm] -> Input 29, Output 30
--> QuantAdd6 -> Input 30 + 28, Output 31
--> ReLU6 -> Input 31, Output 32
--> SparseConvolutionQunat18[subm] -> Input 32, Output 33
--> SparseConvolutionQunat19[subm] -> Input 33, Output 34
--> QuantAdd7 -> Input 34 + 32, Output 35
--> ReLU7 -> Input 35, Output 36
--> SparseConvolutionQunat20[conv] -> Input 36, Output 37
--> ToDense0[[180, 180, 2]][[1, 128, 180, 180, 2]] -> Input 37, Output 38
--> Permute0[(0, 1, 4, 2, 3)][[1, 128, 2, 180, 180]] -> Input 38, Output 39
--> Reshape0[(1, 256, 180, 180)] -> Input 39, Output 40
Tracing done!