目录
- 前言
- 1. 稀疏卷积
- 2. Sparse Convolution Model
- 2.1 输入数据模型
- 2.2 卷积核
- 2.3 输出的定义
- 2.4 计算流程
- 2.4.1 构建 hash table
- 2.4.2 构建 Rulebook
- 2.4.3 在GPU上计算Pipeline
- 2.5 Summary
- 3. SCN导出
- 3.1 实现trace
- 3.2 导出onnx
- 3.3 CenterPoint SCN导出
- 3.4 执行图的构建
- 3.5 onnx解析并创建执行图
- 4. 补充知识
- 4.1 体素化相关
- 4.2 trace相关
- 结语
- 参考
前言
杜老师推出的 复杂onnx解决方案（以sparseconv为例） 课程，通过本次课程学习稀疏卷积以及 SCN 模型的导出细节，记录下个人学习笔记，仅供自己参考。
1. 稀疏卷积
我们先来了解下稀疏卷积和对应的一些基础知识
正常卷积如 图1 所示，需要 9 次乘，而稀疏卷积仅需要 3 次乘
我们在 图1-1 中观察到矩阵中存在大量 0 元素，因此可以采用稀疏化储存表示
- features = [1 2 5]，代表非 0 元素值，一般用二维数组表示 n x 1
- indices = [[0, 0], [1, 1], [2, 2]]，表示非 0 元素值的索引，一般用二维数组表示 n x 2
卷积核只需要与非 0 元素值进行运算即可，因此只需要 3 次乘
那为什么需要稀疏卷积呢？在什么情况下会存在大量的零元素呢？
稀疏主要用在点云数据上，假设现有 100 个点，每个点对应有 k 个特征，则可稀疏表示为
- features.shape = 100 x k
- indices.shape = 100 x 4 这里 4 分别表示 [batch, x, y, z]，一个点有 xyz 在 voxel grid 中的坐标，以及属于哪个 batch
稀疏卷积就是标准卷积的稀疏版本，其输入是稀疏的。
对于稀疏卷积有两种形式：(copy自here)
一种是 Spatially Sparse Convolution，在 spconv 中为 SparseConv3D。就像普通的卷积一样，只要卷积核 kernel 覆盖了一个非零输入点，就会计算对应的输出。对应论文 SECOND:Sparsely Embedded Convolutional Detection
另一种是 Submanifold Sparse Convolution，在 spconv 中为 SubMConv3D。只有当卷积核 kernel 中心覆盖一个非零输入点时，卷积输出才会被计算。对应论文 3D Sematic Segmentation with Submanifold Sparse Convolutional Networks
更多细节可查看下面的第 2.3 小节内容
2. Sparse Convolution Model
卷积神经网络已被证明对二维图像信号处理非常有效。然而，对于三维点云信号，额外的维度 Z 大大增加了计算量。另一方面，与普通图像不同，三维点云的大部分体素是空的，这使得三维体素中的点云数据往往是稀疏的信号。
问题在于我们是否只用稀疏的数据有效地计算卷积，而不是扫描所有的图像像素或空间体素。
很直观的一个想法是，常规的图像信号被存储为矩阵或张量，而相应的卷积计算为密集矩阵的乘法。稀疏信号通常可以被表示为数据列表和索引列表，我们可以采用特殊的卷积模式，利用稀疏性号特殊的表示的优势。
简单来说，传统卷积使用 FFT 或者 im2col 来构建计算管道，而稀疏卷积会收集所有关于卷积核元素的原子操作，并将它们作为计算指令保存在 Rulebook 中
下面我们通过一个例子来解释稀疏卷积的工作原理。
以下内容均 Copy 自 https://towardsdatascience.com/how-does-sparse-convolution-work-3257a0a8fd1
2.1 输入数据模型
为了一步步解释稀疏卷积的概念，让它更容易理解，我们先以二维图像处理为例。由于稀疏信号是采用数据列表和索引列表来表示的，因此二维和三维稀疏信号并没有本质的区别
如 图2-1 所示，我们现在有一个 3 通道的 5x5 大小的图像。除了 P1 和 P2 两个点之外，其余所有像素均为 (0,0,0)。P1 和 P2 这种非零元素也称为 active input sites
在密集形式下，输入张量的 shape 为 [1x3x5x5]，对应的维度分别为 NCHW；在稀疏形式下，数据列表为 [[0.1, 0.1, 0.1], [0.2, 0.2, 0.2]]，索引列表为 [[1, 2], [2, 3]] (采用 YX 顺序)
2.2 卷积核
稀疏卷积的卷积核和传统卷积核相同。图2-2 中的 kernel size 为 3x3。深色和浅色分别代表 2 个 filter，在这个例子中我们使用如下的卷积参数：
conv2D(kernel_size=3, out_channels=2, stride=1, padding=0)
2.3 输出的定义
稀疏卷积的输出与传统卷积有很大不同。稀疏卷积有 2 种输出定义。
一种是 regular output 定义，就像普通的卷积一样，只要卷积核 kernel 覆盖了一个非零输入点，就会计算对应的输出。
另一种被称为 submanifold output 定义，只有当卷积核 kernel 中心覆盖一个非零输入点时，卷积输出才会被计算。
图2-3 说明了这两种输出的区别，A1 代表 active output sites 也就是 P1 的卷积结果；同样地，A2 也代表 active output sites 是由 P2 计算出来的。A1A2 也代表 active output sites 是 P1 和 P2 的输出之和。深色和浅色代表着不同的输出通道。
因此，在密集形式下，输出张量的形状为 [1x2x3x3] 其维度顺序为 NCHW；在稀疏形式下，输出是由两个列表构成，一个是数据列表，一个是索引列表，这与稀疏形式下的输入的表达是类似的。
2.4 计算流程
传统卷积通常使用 im2col 将卷积重写并将卷积视为密集矩阵乘法问题。然而，稀疏卷积使用一个 Rulebook 来计算所有的原子操作，而不是使用 im2col。
2.4.1 构建 hash table
第一步是来构建 hash tables
完整的流程如 图2-4 所示，在 图2-4 中输入哈希表存储了所有 active input sites ，然后，我们估计了所有可能的 active output sites，并考虑两种输出定义 (Sparse Output、Submanifold Output) 中的某一种去计算输出，最后使用输出哈希表来记录所有的 active output sites
值得注意的是，在 图2-4 中对于键值的描述并不是那么清楚， v v v 更像一个哈希键，而 k e y key key 更像一个哈希值，它们都没有重复的元素。因此，具体哪一个应该是键，哪一个应该是值，取决于你的程序。
接下来我们一步步来介绍构建过程，先来看下 P1 是怎么进行卷积操作的：
但是，并不是每次 kernel 在卷积过程中都可以碰到 P1，从第七次开始，输出矩阵就不再发生变化了。
得到 P1 的输出后我们来记录每个元素的位置
上面的只是操作 P1，对于 P2 也是同样的操作，如下图所示：
最后把 P1，P2 的结果结合起来(主要是消除掉重复元素)，得到一张位置表，编号得到 output hash table
2.4.2 构建 Rulebook
第二步是建立 Rulebook，这是稀疏卷积的关键部分。Rulebook 的目的类似于 im2col，它将卷积从数学形式转换为有效的可编程形式。但与 im2col 不同的是，Rulebook 收集了卷积中所有涉及的原子操作，然后将它们关联到相应的 kernel 元素上。
图2-8 演示了如何构建 Rulebook， P i n P_{in} Pin 中包含着输入索引，在这个示例中，我们在位置 (2,1) 和 (3,2) 有两个非零元素， P o u t P_{out} Pout 中有相应的输出索引，接下来，我们会收集卷积计算过程中的原子操作(atomic operator)，即把卷积过程看成是许多原子操作。最后，我们将所有的原子操作记录在 Rulebook 中， 在 图2-8 的 Rulebook 中，第一列是 kernel element 的索引，第二列是一个计数器，关于这个 kernel element 涉及多少个原子操作。第三列和第四列是关于这个原子操作中输入哈希表的索引和输出哈希表的索引。
( i , j ) (i,j) (i,j) 具体是如何生成的呢？很简单，我们在做卷积的时候对应的非零元素 P 在 single channel kernel template 的位置就是对应的 ( i , j ) (i,j) (i,j)，比如在 P1 第一次卷积时，P1 对应的位置是 (+1,0)，因此其对应的 ( i , j ) (i,j) (i,j) 就是 (+1,0)
2.4.3 在GPU上计算Pipeline
最后是在程序中实现稀疏卷积，过程如下：
图2-10 是计算稀疏卷积的示例，初看可能觉得复杂，细看你会发现也不简单😂。如 图2-10 所示，我们在计算稀疏卷积时没有采用传统卷积的滑动窗口方法，而是根据 Rulebook 计算所有的原子操作。
在 图2-10 中，红色和蓝色的箭头表示两个计算实例。红色箭头处理 kernel 元素 (-1,-1) 的第一个原子操作，从 Rulebook 中我们可知这个原子操作的输入位置是 P1(2,1)，输出位置是(2,1)，整个原子操作的流程如 图2-11 所示。
同样地，蓝色箭头表示另一个原子操作，它共享相同的输出。红色箭头实例和蓝色箭头实例的结果可以加在一起。
然而，在实际应用中，神经网络计算的是具有激活函数的线性变换之和，如 f ( ∑ i w i x i + b ) f\left(\sum_{i} w_i x_i+b\right) f(∑iwixi+b) ，因此我们其实并不需要输出的哈希表，只需将所有原子操作的结果相加即可
值得注意的是，关于 Rulebook 的计算在 GPU 中可实现并行。
2.5 Summary
稀疏卷积是相当有效的，因为我们不需要扫描所有的像素或空间体素。我们只计算非零元素的卷积。我们使用 Rulebook 而不是 im2col 将稀疏卷积转换成一个紧凑的线性变换问题，或者说一个紧凑的矩阵乘法问题。一个额外的计算成本是建立 Rulebook，幸运的是，这个构建工作可以在 GPU 上并行处理
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 实现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)
运行效果输入如下：
在上面的示例代码中，我们通过替换特定函数的实现，实现了钩子函数来挂钩到自定义函数中。具体来说，我们定义了一个名为 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)
运行效果如下：
在上述示例代码中，我们定义了一个 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)
运行效果如下：
上述示例代码通过引用 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 文件组成如下图所示：
- 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 格式的文件。
运行效果如下：
导出的 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)
可以看到 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 时，里面再次触发 hookdef 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.4 执行图的构建
前面我们解决了模型的 onnx 导出问题，现在来探讨模型推理问题，由于 tensorRT 这条路走不通，因此只能自行解析 onnx 单独进行推理。
onnx 是静态的计算图，我们需要构建一个执行图(如图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))
运行效果如下：
上述代码演示了一个简化的 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))
运行效果如下：
至此，整个 onnx 解析器已经完成。你应该在 C++ 上复现这个流程就可以实现无 trt 推理了。
4. 补充知识
4.1 体素化相关
我们来补充下体素化的相关知识(from chatGPT)
体素 (voxel) 是"体积元素"的简称，是二维图像中像素的三维等效物。它们代表了三维空间中的最小单位，通常被视觉化为小立方体或立方体。在基于体素的表述中，三维物体或场景被离散成体素的网格，每个体素可以存储各种信息，如颜色、密度或材料属性。
体素化 (voxelization) 是将连续的几何表示 (如多边形网格) 转换为基于体素的表示的过程。它涉及到将物体或场景所占据的三维空间划分为一个有规律的体素网格，并根据物体几何形状的存在与否来确定哪些体素是"填充"或"空"的。
点云体素化 (point cloud voxelization) 是将点云，即空间中的三维点的集合，转换为基于体素的表示的过程。点云是一种流行的数据结构，用于表示从各种传感器获得的三维几何数据，如三维扫描仪、LiDAR传感器或深度相机。
点云的体素化涉及到将点所占据的三维空间离散为一个有规律的体素网格。然后，每个体素可以根据其边界点的存在或不存在而被分类为"占用"或"空"。这个过程有效地将连续和非结构化的点云数据转化为结构化的体素网格表示。(如 图4-1 所示)
将点云体素化有几个有点。它提供了一个紧凑和统一的表示，能够有效地存储、处理和分析三维数据。基于体素的表示非常适合于碰撞检测、表面重建、物体识别和空间推理等任务。体素化也有利于点云数据与其它基于体素的算法和技术的整合。
4.2 trace相关
PyTorch 中
trace是什么呢？(from chatGPT)在 PyTorch 中，
trace是一种用于从给定的输出示例生成模型的图形表示的方法。它将模型的前向传播过程转化为静态图形，并记录计算图中的每个操作和边缘输入。简单来说，
trace是用于生成 ONNX 模型的过程，它会跟随模型的前向传播过程，并记录下每个节点及其对应的权重信息。通过trace，PyTorch 可以将动态图形转换为静态图形，其中每个节点表示模型的操作，例如张量运算、激活函数等。在跟踪过程中，每个节点的权重信息也会被记录下来，以便在导出的 ONNX 模型中正确重现模型的计算。
PyTorch 中
jit trace是什么呢？(from chatGPT)
jit trace是 PyTorch 的一种工具，用于将模型的前向传播过程转化为静态计算图，并支持导出 ONNX 格式简单来说就是 PyTorch 官方实现的一种 trace 方法，可以支持导出 ONNX 格式
jit代表 “Just-in-Time”，是一种即时编译的技术，它可以将 PyTorch 模型的动态计算图转换为静态计算图，以提高执行效率。jit模块提供了多种转换和优化工具，包括trace函数在 PyTorch 中，可以使用
torch.jit.trace函数来执行jit trace。该函数接受一个模型和输入示例，并返回一个经过跟踪的torch.jit.ScriptModule对象，其中包含了模型的静态计算图表示这个
jit trace的功能与trace类似，但它是 PyTorch 官方实现的，专门用于支持将模型转换为静态计算图，并导出为 ONNX 格式。通过jit trace，可以方便地将模型转换为 ONNX 模型，并在其它框架或设备上进行推理和部署。
动态计算图和静态计算图又是什么呢？(from chatGPT)
在动态计算图中，计算图是在运行时动态构建的。这意味着在每次前向传播过程中，计算图会根据输出数据的形状和值来动态生成。这种动态性使得在模型的构建和调试过程中更加灵活，可以使用条件语句、循环和可变控制流等动态特性。PyTorch 就是一个基于动态计算图的框架，它允许用户以一种更直观和灵活的方式定义和修改模型。
相比之下，静态计算图在模型定义阶段就被预先定义和固定，然后在执行阶段进行计算。例如在 TensorFlow 中，计算图需要在构建阶段进行定义，包括网络结构、操作和张量之间的连接关系。值得注意的是，最近的版本中 TensorFlow 也引入了 Eager Execution 的概念，它允许用户以动态计算图的方式进行模型的定义和执行。ONNX 模型本身也可以被视为一种静态计算图的表示形式。
回到 SPConv 的 onnx 导出问题，由于 PyTorch 官方实现的 jit trace 并不支持 SPConv 的 ONNX 导出，因此需要自己去实现 trace 并导出 ONNX
自定义实现
trace时，有几个注意事项：跟踪静态图：确保在模型的前向传播过程中，所有的操作都是静态的，即不依赖于动态控制流（如循环或条件语句）。这是因为
trace需要将模型表示为静态图形，以便正确地捕捉和记录计算。输入示例的准备：为了进行跟踪，需要提供输入示例，以便 PyTorch 了解输入张量的形状和类型。输入示例应与模型的实际输入具有相同的形状和类型。如果模型的输入是多维的，可能需要提供一批示例输入。
不可变参数：在跟踪期间，模型的参数应保持不变。确保在跟踪期间不要修改模型的参数，以免引起意外的结果。
异常处理：跟踪期间可能会遇到无法跟踪的操作或数据类型。在自定义实现
trace时，要注意处理这些异常情况，以便正确处理模型的各个部分。
PyTorch 中
register_forward_hook是什么？(from chatGPT)在 PyTorch 中，
register_forward_hook是一种用于模型的前向传播过程注册钩子（hook）的方法。钩子是一种在模型的前向传播过程中插入自定义操作的机制。
register_forward_hook允许用户在模型的每个模块的前向传播过程中注册一个函数，该函数将在模块的输出被计算出来之后被调用。
register_forward_hook的主要用途包括：监听模型的中间层输出：通过注册前向钩子，可以获取模型中间层的输出。这对于可视化模型的中间特征或进行模型分析和解释非常有用。
修改模型的输出：可以使用前向钩子来修改模型的输出，例如在模型输出之后应用后处理操作或添加自定义损失函数。
收集模型的统计信息：通过注册前向钩子，可以在每个批次或每个样本级别上收集模型的统计信息，如激活值分布、梯度值等。(例如量化中敏感层的分析？🤔)
当使用 register_forward_hook 时，需要提供一个用于处理前向传播输出的函数，并将其注册到模型的指定模块上。
以下是一个简单的示例代码，展示了如何使用 register_forward_hook 来获取模型中间层的输出：
import torch
import torch.nn as nn
# 定义模型
class Model(nn.Module):
def __init__(self):
super().__init__()
self.fc1 = nn.Linear(784, 256)
self.fc2 = nn.Linear(256, 64)
self.fc3 = nn.Linear(64, 10)
def forward(self, x):
x = torch.relu(self.fc1(x))
x = torch.relu(self.fc2(x))
x = self.fc3(x)
return x
# 定义前向钩子函数
def hook_fn(moduel, input, output):
print(f"模块名：{moduel.__class__.__name__}")
print(f"输入：{input}")
print(f"输出：{output}")
# 创建模型实例
model = Model()
# 注册前向钩子函数
model.fc2.register_forward_hook(hook_fn)
# 创建随机输入
input = torch.randn(1, 784)
# 前向传播
output = model(input)
运行效果输出如下：
在上面的示例代码中，我们定义了一个包含三个线性层的简单模型 Model。然后，我们定义了一个 hook_fn 函数作为前向钩子，用于处理模块的输出。我们将前向钩子注册到模型的 fc2 模块上。最后，我们创建了一个随机输入并进行前向传播。
当前向传播到 fc2 模块时，前向钩子函数 hook_fn 将被调用，打印该模块的名称、输入和输出。通过这种方式，我们可以获取模型中间层的输出，并进行相应的处理或分析。
需要注意的是，可以在模型的不同模块上注册多个前向钩子，以便在多个位置获取输出。
在 register_forward_hook 中注册的前向钩子函数可以具有不同的输入参数，它可以接收 module、input 和 output 三个参数中的一个或多个
module：表示当前调用钩子的模块。可以使用module参数来获取模块的相关信息，如模块的名称、类型等input：表示模块的输入。对于具有多个输入的模块，input参数是一个元组，其中每个元素对应一个张量output：表示模块的输出。对于具有多个输出的模块，output参数是一个元组，其中每个元素对应一个输出张量
结语
本篇博客从课程出发，了解了稀疏卷积的相关知识，跟随杜老师的脚步学习了以 spconv 为例的复杂 onnx 的解决方案，当需要导出的算子不支持 jit trace 时需要我们自定义实现，通过钩子函数去替换 forward 从而实现一些自定义的操作。同时当导出的 onnx 无法对接 tensorRT 时，需要我们自行解析 onnx，将思维打开，不要局限在 tensorRT、onnxruntime 这些东西，而是能够自行处理和解决很多具体的事情。
博主能力有限，许多知识并没有掌握，只是分享下相关知识为后续的学习做个铺垫。博主在这里只做了简单分享，很多具体的实现并没有去做，比如没有去配置 spconv 的环境进行验证测试，单纯从跟随杜老师从原理的角度去做了简单的分析。更多细节需要给位看官自行去了解了上述😄。
感谢各位看到最后，创作不易，读后有收获的看官请帮忙点个👍
参考
- 复杂onnx解决方案（以sparseconv为例）
- 3D稀疏卷积粗略理解：Submanifold Sparse Convolution和Spatial Sparse Convolution以及SECOND网络理解
- 通俗易懂的解释Sparse Convolution过程
- How does sparse convolution work?
- SECOND:Sparsely Embedded Convolutional Detection
- 3D Sematic Segmentation with Submanifold Sparse Convolutional Networks
- 通用矩阵乘(GEMM)优化算法
- https://github.com/traveller59/spconv
- https://github.com/tianweiy/CenterPoint/blob/master/det3d/models/backbones/scn.py
