目录
- 前言
- 0. 简述
- 1. 使用spconv进行SCN的推理测试
- 2. 导出onnx
- 3. 补充-装饰器+钩子函数
- 总结
- 下载链接
- 参考
前言
自动驾驶之心推出的 《CUDA与TensorRT部署实战课程》，链接。记录下个人学习笔记，仅供自己参考
本次课程我们来学习下课程第八章——实战：CUDA-BEVFusion部署分析，一起来学习导出带有 spconv 的 SCN 网络的 onnx
Note：之前在学习杜老师的课程中有简单记录过 Sparse Convolution 的一些基础知识，感兴趣的可以看下：复杂onnx解决方案（以sparseconv为例）
课程大纲可以看下面的思维导图
0. 简述
本小节目标：学习利用 hook 截取 spconv 的 forward，从而自定义 onnx 算子导出 onnx 的方法
今天给大家讲解第八章第 3 小节，学习导出带有 spconv 的 SCN 网络的 onnx，这个小节我们先跟着 NVIDIA 官方提供的 3D Sparse Convolution 库，学习怎么去调用它，接口是什么样子
我们先假设 onnx 已经导出好了，读取导出好的 onnx 生成对应的 engine 引擎完成前向推理，我们先来完成这个流程，onnx 导出我们稍后再看
这部分 NVIVIDA 其实开源在 https://github.com/NVIDIA-AI-IOT/Lidar_AI_Solution/tree/master/libraries/3DSparseConvolution，但值得注意的是它整个推理框架虽然是开源了，但比较核心的地方比如 spconv 里面是怎么加速的它其实并没有开源
我们主要是通过 NVIDIA 提供的方案一起去学习一下它的推理框架是怎么做的以及 spconv 的接口是如何去使用的，这是我们需要做的第一件事
第二件事就是我们需要学习 spconv 的 onnx 是怎么导出的，我们上节课也讲过 spconv 的 onnx 稍微有点特殊，因为我们需要用 hook 去截取 spconv 的 forward，之后重定位 spconv 的forward，接着去创建一个 onnx 自定义节点，通过这一系列的操作完成 spconv 的导出
1. 使用spconv进行SCN的推理测试
我们先来看一下 NVIDIA 官方提供的 3D Sparse Convolution 的 README 文档，它位于 LiDAR_AI_Soultion 项目下的 libraries/3DSparseConvolution，README 中提到 3DSparseConvolution 这个库提供了一个 int8/fp16 精度的 3D 稀疏卷积网络的推理引擎，如下图所示：
我们先来将它跑通，README 中提到该项目的运行过程包括以下三部分：
- 1. 从 https://github.com/tianweiy/CenterPoint 下载并配置 CenterPoint 环境
- 2. 导出 SCN ONNX
- 3. 编译和运行
这里有个小技巧，那就是我们可以先暂时跳过 SCN ONNX 的导出，因为在下载 LiDAR_AI_Solution 这个 repo 时官方其实提供了导出好的 CenterPoint SCN 的 ONNX 供我们测试，具体位置在 3DSparseConvolution/workspace/centerpoint 文件夹下，因此我们可以跳过第一部分和第二部分，直接利用官方提供的导出好的 SCN ONNX 来编译运行看能否通过
因此我们先假设这个 ONNX 导出好了，我们来直接读取这个 ONNX 看看它是怎么做前向推理的，指令如下：
$ sudo apt-get install libprotobuf-dev=3.6.1*
$ cd path/to/3DSparseConvolution
$ make fp16 -j
该库依赖于 protobuf，因此你需要先安装 protobuf，关于 protobuf 库的安装博主在 八. 实战：CUDA-BEVFusion部署分析-环境搭建 中也提供了 apt 和源码两种安装方式，这边不再赘述。接着我们就可以直接执行 make 指令来编译了，输出如下图所示：
我们简单看一下编译过程中它都做了些什么，从日志中我们可以看到它先做了一个 Parse node conv0 [Sparseconvolution]，把这个 Node 节点 Parse 完之后再做一个 add，也就是把 Parse 完之后的信息读取出来加入到 conv0 里面
那 conv0 这个东西我们从这里面可以看到它其实是一个自定义的节点，它里面参数是 submanifold，也就是我们说的 3D Sparse Convolution 的一种形式，同样依此类推后面还有 conv1、conv2 等等，这一系列操作都是这么做的
其次我们再来关注下其精度，从日志中可以看出输入输出部分是跑的 FP16 精度，而中间的 spconv 部分跑的是 INT8 精度，因此 SCN 整体推理框架的精度是 FP16+INT8
我们再来看看整个网络 forward 过程中 tensor 维度的变化，我们可以看到输入维度是 41x1440x1440，1440x1440 代表着 3D 点云经过体素化之后得到的一个坐标系，41 代表它的通道数
我们可以看到它整个过程中其实就是不断的进行稀疏卷积，然后做 downsample 变成 720x720 之后再做一个稀疏卷积降维到 360x360 之后再通过稀疏卷积降维到 180x180。降到 180x180 之后我们再做一个 scatter，那这个操作就是将维度扩充到 1x128x2x180x180 的维度，之后再做 reshape，reshape 到 1x256x180x180 的 tensor，这个就是 SCN 网络最终输出 tensor 的一个维度了
那 1x255x180x180 代表的意义是它生成的 BEV Grid 的大小是 180x180，每一个 Grid 上有 256 维的特征，那这 256 维的特征是从点云那边直接学习得到的，以上就是 SCN 网络 forward 的整个流程
OK，下面我们来进代码看看内部是怎么实现的，我们先从 main.cpp 看起，代码入口如下：
int main(int argc, char** argv) {
const char* cmd = "fp16";
if (argc > 1) cmd = argv[1];
cudaStream_t stream = nullptr;
checkRuntime(cudaStreamCreate(&stream));
if (strcmp(cmd, "memint8") == 0) do_memory_usage_test(spconv::Precision::Int8, stream);
if (strcmp(cmd, "memfp16") == 0) do_memory_usage_test(spconv::Precision::Float16, stream);
if (strcmp(cmd, "int8") == 0) do_simple_run(spconv::Precision::Int8, stream);
if (strcmp(cmd, "fp16") == 0) do_simple_run(spconv::Precision::Float16, stream);
checkRuntime(cudaStreamDestroy(stream));
return 0;
}
当我们执行 make fp16 -j 编译运行时，代码首先会调用 do_memory_usage_test 函数来测试稀疏卷积操作的内存使用情况，接着我们会调用 do_simple_run 函数来执行稀疏卷积，我们重点来看下 do_simple_run 这个函数看它内部的稀疏卷积是怎么执行的，其代码实现如下：
void do_simple_run(spconv::Precision precision, cudaStream_t stream) {
spconv::set_verbose(true);
auto task = load_task("centerpointZYX", precision);
// auto task = load_task("bevfusionZYX", precision);
// auto task = load_task("bevfusionXYZ", precision);
task.engine->input(0)->set_data(task.features.shape, spconv::DataType::Float16, task.features.ptr(),
task.indices.shape, spconv::DataType::Int32, task.indices.ptr(),
task.grid_size);
task.engine->forward(stream);
auto out_features = task.engine->output(0)->features();
auto grid_size = task.engine->output(0)->grid_size();
printf("🙌 Output.shape: %s\n", spconv::format_shape(out_features.shape).c_str());
out_features.save(task.save_dense, stream);
task.engine.reset();
print_done(task.compare_cmd);
}
首先我们先做了一个 load_task 加载了一个 centerpointZYX 的任务，load_task 的代码实现如下：
Task load_task(const string& name, spconv::Precision precision) {
Task task;
task.name = name;
if (name == "bevfusionXYZ") {
task.engine = spconv::load_engine_from_onnx("bevfusion/bevfusion.scn.xyz.onnx", precision);
task.features = spconv::Tensor::load("bevfusion/infer.xyz.voxels");
task.indices = spconv::Tensor::load("bevfusion/infer.xyz.coors");
task.grid_size = {1440, 1440, 41};
task.order = IndiceOrder::XYZ;
task.save_dense = "bevfusion/output.xyz.dense";
task.compare_cmd =
"python tool/compare.py workspace/bevfusion/infer.xyz.dense "
"workspace/bevfusion/output.xyz.dense --detail";
} else if (name == "bevfusionZYX") {
task.engine = spconv::load_engine_from_onnx("bevfusion/bevfusion.scn.zyx.onnx", precision);
task.features = spconv::Tensor::load("bevfusion/infer.zyx.voxels");
task.indices = spconv::Tensor::load("bevfusion/infer.zyx.coors");
task.grid_size = {41, 1440, 1440};
task.order = IndiceOrder::ZYX;
task.save_dense = "bevfusion/output.zyx.dense";
task.compare_cmd =
"python tool/compare.py workspace/bevfusion/infer.zyx.dense "
"workspace/bevfusion/output.zyx.dense --detail";
} else if (name == "centerpointZYX") {
task.engine = spconv::load_engine_from_onnx("centerpoint/centerpoint.scn.PTQ.onnx", precision);
task.features = spconv::Tensor::load("centerpoint/in_features.torch.fp16.tensor");
task.indices = spconv::Tensor::load("centerpoint/in_indices_zyx.torch.int32.tensor");
task.grid_size = {41, 1440, 1440};
task.order = IndiceOrder::ZYX;
task.save_dense = "centerpoint/output.zyx.dense";
task.compare_cmd =
"python tool/compare.py workspace/centerpoint/out_dense.torch.fp16.tensor "
"workspace/centerpoint/output.zyx.dense "
"--detail";
} else {
Assertf(false, "Unsupport task name: %s", name.c_str());
}
return task;
}
由于我们 name 等于 centerpointZYX，因此这里用它做了一系列初始化，先通过读取 onnx 创建了一个 engine，我们来看下它具体是怎么创建的，也就是 load_engine_from_onnx 具体是怎么做的，其代码如下：
std::shared_ptr<Engine> load_engine_from_onnx(const std::string& onnx_file, Precision precision, void* stream, bool mark_all_output){
onnx::ModelProto model;
std::fstream fin(onnx_file, std::ios::binary | std::ios::in);
if (!model.ParseFromIstream(&fin)) {
LOGV("Parse onnx failed: %s", onnx_file.c_str());
return nullptr;
}
auto builder = spconv::create_engine_builder();
auto graph = model.graph();
std::unordered_map<std::string, spconv::ITensor*> tensor_map_by_name;
for (int i = 0; i < graph.input_size(); ++i) {
auto name = graph.input(i).name();
tensor_map_by_name[name] = builder->push_input(name);
}
std::vector<spconv::ITensor*> collect_outputs;
for (int i = 0; i < model.graph().node_size(); ++i) {
auto& node = model.graph().node(i);
if (node.op_type() == "SparseConvolution") {
auto x = tensor_map_by_name[node.input(0)];
auto weight = get_initializer_data(graph, node.input(1));
auto bias = get_initializer_data(graph, node.input(2));
auto weight_dynamic_ranges_proto = get_attribute(node, "weight_dynamic_ranges");
auto weight_dynamic_ranges =
std::vector<float>(weight_dynamic_ranges_proto.floats().begin(), weight_dynamic_ranges_proto.floats().end());
auto n = builder->push_sparse_conv(
node.name(), x,
weight.data, weight.shape,
weight_dynamic_ranges,
bias.data, bias.shape,
get_attribute(node, "activation").s(),
get_attribute_as_intarray(node, "kernel_size"),
get_attribute_as_intarray(node, "stride"),
get_attribute_as_intarray(node, "padding"),
get_attribute_as_intarray(node, "dilation"),
get_attribute(node, "input_dynamic_range").f(),
get_attribute(node, "subm").i(),
get_attribute(node, "output_bound").i(),
get_attribute(node, "rulebook").s(),
get_attribute(node, "precision").s() == "int8" ? Precision::Int8 : Precision::Float16,
get_attribute(node, "output_precision").s() == "int8" ? Precision::Int8 : Precision::Float16,
node.output(0)
);
if(mark_all_output){
collect_outputs.push_back(n->output(0));
}
tensor_map_by_name[node.output(0)] = n->output(0);
} else if (node.op_type() == "Add" || node.op_type() == "QuantAdd") {
auto a = tensor_map_by_name[node.input(0)];
auto b = tensor_map_by_name[node.input(1)];
auto n = builder->push_add(
node.name(),
a, b,
get_attribute(node, "input0_dynamic_range").f(),
get_attribute(node, "input1_dynamic_range").f(),
node.output(0),
get_attribute(node, "precision").s() == "int8" ? Precision::Int8 : Precision::Float16,
get_attribute(node, "output_precision").s() == "int8" ? Precision::Int8 : Precision::Float16
);
tensor_map_by_name[node.output(0)] = n->output(0);
} else if (node.op_type() == "Relu") {
auto x = tensor_map_by_name[node.input(0)];
auto n = builder->push_relu(node.name(), x, node.output(0));
tensor_map_by_name[node.output(0)] = n->output(0);
} else if (node.op_type() == "ScatterDense") {
auto x = tensor_map_by_name[node.input(0)];
auto input_spatial_shape = get_attribute_as_intarray(node, "input_spatial_shape");
auto output_shape = get_attribute_as_intarray(node, "output_shape");
auto format = get_attribute(node, "format").s();
auto n = builder->push_dense(node.name(), x, format, node.output(0), input_spatial_shape, output_shape);
tensor_map_by_name[node.output(0)] = n->output(0);
} else if (node.op_type() == "Reshape") {
auto x = tensor_map_by_name[node.input(0)];
auto dims = get_attribute(node, "dims");
std::vector<int64_t> shape(dims.ints().begin(), dims.ints().end());
auto n = builder->push_reshape(node.name(), x, shape, node.output(0));
tensor_map_by_name[node.output(0)] = n->output(0);
} else if (node.op_type() == "Transpose") {
auto x = tensor_map_by_name[node.input(0)];
auto dims = get_attribute(node, "dims");
std::vector<int64_t> shape(dims.ints().begin(), dims.ints().end());
auto n = builder->push_transpose(node.name(), x, shape, node.output(0));
tensor_map_by_name[node.output(0)] = n->output(0);
} else {
printf("Unsupport operator [%s]\b", node.op_type().c_str());
return nullptr;
}
}
for (int i = 0; i < graph.output_size(); ++i) {
auto name = graph.output(i).name();
collect_outputs.push_back(tensor_map_by_name[name]);
}
for (int i = 0; i < collect_outputs.size(); ++i) {
builder->push_output(collect_outputs[i]);
}
return builder->build(precision, stream);
}
我们从名字也能看出来，它的功能是读取一个 onnx 生成一个 engine，从代码中能看出先利用 protobuf 的接口函数先解析 onnx 文件，接着创建了一个 builder，我们来看一下这个 builder 是一个什么类，其代码如下：
class EngineBuilder{
public:
Exported virtual ITensor* push_input(const std::string& name) = 0;
Exported virtual INode* push_add(
const std::string& name,
ITensor* a,
ITensor* b,
float a_dynamic_range,
float b_dynamic_range,
const std::string& output_name,
Precision precision, Precision output_precision) = 0;
Exported virtual INode* push_relu(
const std::string& name,
ITensor* x,
const std::string& output_name) = 0;
Exported virtual INode* push_dense(
const std::string& name, ITensor* x,
const std::string& format,
const std::string& output_name,
const std::vector<int>& input_spatial_shape,
const std::vector<int>& output_shape) = 0;
Exported virtual INode* push_reshape(
const std::string& name, ITensor* x,
const std::vector<int64_t>& shape,
const std::string& output_name) = 0;
Exported virtual INode* push_transpose(
const std::string& name, ITensor* x,
const std::vector<int64_t>& dims,
const std::string& output_name) = 0;
Exported virtual INode* push_sparse_conv(
const std::string& name,
ITensor* x,
const std::vector<unsigned short>& weight,
const std::vector<int>& weight_shape,
const std::vector<float>& weight_dynamic_ranges,
const std::vector<unsigned short>& bias,
const std::vector<int>& bias_shape,
const std::string& activation,
const std::vector<int>& kernel_size,
const std::vector<int>& stride,
const std::vector<int>& padding,
const std::vector<int>& dilation,
float input_dynamic_range,
bool submanifold,
int max_output_points,
const std::string& rulebook,
Precision precision,
Precision output_precision,
const std::string& output_name) = 0;
Exported virtual void push_output(ITensor* value) = 0;
// build engine
Exported virtual std::shared_ptr<Engine> build(Precision precision, void* stream = nullptr) = 0;
};
/**
* To build a engine.
*/
Exported std::shared_ptr<EngineBuilder> create_engine_builder();
我们可以看到 builder 是一个 EngineBuilder 类，那这个名字其实和 TensorRT 里面的 IBuilder 比较像，但是这两个并不是一个东西，IBuilder 是 TensorRT 的概念，这个 EngineBuilder 是 NVIDIA 官方自己重新写的一个 builder，它没有集成任何东西是官方自己构建的一个 builder，那既然是 builder 我们就知道它肯定是用于来创建 engine 的
OK，我们再回到 load_engine_from_onnx 函数中，创建完 builder 之后我们再拿到模型的 graph 图结构，然后我们会遍历 graph 中的所有 node，针对不同的 node 进行不同的操作，那这些 node 包括 SparseConvolution 也就是我们说的稀疏卷积，还有 Add、QuantAdd、Relu、ScatterDense 等等，
如果 Node 节点是 SparseConvolution 稀疏卷积，我们会先通过 get_initializer_data 函数拿到它的 weights 和 bias 数据，接着通过 push_sparse_conv 函数把稀疏卷积的权重和偏置数据，对应的属性等等添加到 builder 中，那具体的 push_sparse_conv 内部是如何实现的其实我们并不知道，官方只提供了一个 engine.hpp 接口文件，那怎么实现的里面的 engine.cpp 或者 engine.cu 它其实并没有开源
其实我们只要知道接口就知道怎么用了，那通过名字也知道 push_sparse_conv 也就是创建了一个 spconv 的节点，我们往这个节点里面放很多信息，比如名字、权重、动态范围等等各种属性，那这样就把 spconv 的一个节点创建了，那对于 onnx 中的每一个节点我们都可以先从 onnx 中 parse 信息，之后把信息 push 到 builder 中去，通过 builder 创建一个节点，那就是这么一个流程，后面的 add、relu、reshape 等节点都可以这么做
所以 load_engine_from_onnx 函数执行完之后我们其实就得到了对应 onnx 的 engine 推理引擎了，那值得注意的是这个 engine 引擎它跟 TensorRT 的 engine 还不一样，这是它自己写的一个 engine
OK，我们再回到 load_task 函数中，加载完 engine 之后，我们会 load 两个 Tensor 一个是 feature 也就是输入到网络中的 3D 点云特征，一个是 indices 也就是输入到网络中的 3D 点云的位置索引，其实也就是加载了 SCN 网络前向推理需要的两个输入 Tensor，接着设置了一些 grid_size、order 等参数
OK，我们再回到 do_simple_run 函数中，通过 load_task 加载了 engine 之后我们通过 set_data 设置了输入的数据，接着调用 forward 进行了前向传播，代码如下所示：
/**
Engine types for sparse convolution
**/
class Engine {
public:
/**
Inference function for sparse convolution
features_shape: The shape of the input feature matrix, it must be two elements.
features_dtype: The data type of the input feature matrix, it must be Float16 now.
features_data: The data pointer of the input feature matrix
indices_shape: The shape of the input indices matrix, it must be two elements[n, 4]
indices_dtype: The data type of the input indices matrix, it must be Int32 now.
indices_data: The data pointer of the input indices matrix
batch: The batch size of the input, it must be 1 now.
grid_size: The grid size of the input data, For example: 41,1440,1440 or 1440,1440,41
stream: Which stream is expected to enqueue the inference.
**/
Exported virtual void forward(void* stream = nullptr) = 0;
Exported virtual size_t num_input() const = 0;
Exported virtual SparseDTensor* input(unsigned int index) = 0;
Exported virtual size_t num_output() const = 0;
Exported virtual SparseDTensor* output(unsigned int index) = 0;
};
那 forward 部分也是没有开源的，我们只知道它有这个接口，我们并不知道 forward 里面具体干了什么，但是我们可以猜一下，它里面的 forward 可能跟 TensorRT 里面的 forward 比较像，只是这里的 forward 是根据它自己创建的节点一个一个计算的，里面可能会涉及到一些 CUDA 加速计算，那我们这里就跟踪不了
OK，forward 做完之后我们就得到了一个 1x256x180x180 的输出数据，那拿到它之后整个前向推理也就结束了，接着把输出的 1x256x180x180 的 features 保持下来，然后再做一个 reset，整个流程就是这样子
那我们和大家一起简单的过了一遍整体流程，那如果说下次我们要使用它提供的这些接口的时候，其实这么分析就好了，这个是 C++ 部分的推理代码的一个解析
我们最后来看一下精度比较，指令如下：
$ python tool/compare.py workspace/centerpoint/out_dense.torch.fp16.tensor workspace/centerpoint/output.zyx.dense --detail
输出结果如下图所示：
那它会对比 C++ 和 Pytorch 的推理结果，将两个 tensor 进行比较计算二者的差异，我们可以看到大部分数据都是一样的，也有少部分数据存在差异，但是都在误差范围之内，所以是可以忽略的
2. 导出onnx
接下来我们主要看这个 onnx 是怎么导出的，NVIDIA 官方其实开源了相关的导出代码，具体是在 3DSparseConvolution/tool/centerpoint-export 文件夹下
其中 export_tool.py 是核心文件，我们一起来看下其具体的实现，代码如下所示：
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, f"spconv{ilayer}.weight"),
]
if self.bias is not None:
inputs.append(append_initializer(self.bias.data, f"spconv{ilayer}.bias"))
act_type_name = {
tv.gemm.Activation.ReLU : "ReLU",
tv.gemm.Activation.None_ : "None",
tv.gemm.Activation.Sigmoid : "Sigmoid",
tv.gemm.Activation.LeakyReLU : "LeakyReLU"
}
algo_name = {
ConvAlgo.MaskImplicitGemm : "MaskImplicitGemm",
ConvAlgo.MaskSplitImplicitGemm : "MaskSplitImplicitGemm",
ConvAlgo.Native : "Native",
}
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 = act_type_name[self.act_type],
input_shape = x.features.shape,
output_shape = y.features.shape
)
)
它这里写了一个 register_node 的 hook 函数，之后写了一个 symbolic_sparse_convolution 就是各种 symbolic 符号函数，当我们程序在 forward 时候会进行重定位
我们都知道要将 pytorch 模型导出 onnx 的话，我们是需要让这个 pytorch 模型完整走一遍 forward 过程的，因为 forward 的过程中我们才能 trace 里面的各个节点。我们知道如果做 Sparse Convolution 的 forward 的时候，它里面有很多东西我们是 trace 不了的，那么该怎么办呢，我们可以截取，可以做一个 hook，当 trace 到 Sparse Convolution forward 的时候我们让它去执行 symbolic_sparse_convolution 函数里面的操作
在函数里面我们得到 output_size，得到 input 数据，再得到 activation 数据，然后让它去调用 helper.make_node 函数创建一个 node 节点，那我们可以看到这里面有很多参数，那这些参数其实就是 SparseConvolution 它在 C++ 里做前向推理的时候所需要的数据
那我们把这些参数通过 make_node 形式保存下来，以 node 的形式给它添加到 onnx 里面去，那这个就是 spconv 重定位的 forward 的一个实现，那它可以通过这么一个 hook 函数完成这点
@register_node("torch.nn.ReLU.forward")
def symbolic_relu(self, ilayer, y, x):
register_tensor(y)
print(f" --> ReLU{ilayer} -> Input {get_tensor_id(x)}, Output {get_tensor_id(y)}")
nodes.append(
helper.make_node(
"Relu", [get_tensor_id(x)], [get_tensor_id(y)], f"relu{ilayer}"
)
)
@register_node("torch.Tensor.__add__")
def symbolic_add(a, ilayer, y, b):
register_tensor(y)
print(f" --> Add{ilayer} -> Input {get_tensor_id(a)} + {get_tensor_id(b)}, Output {get_tensor_id(y)}")
nodes.append(
helper.make_node(
"Add", [get_tensor_id(a), get_tensor_id(b)], [get_tensor_id(y)], f"add{ilayer}"
)
)
同理，如果说 Sparse Convolution 的 forward 我们可以截取，那么是不是其它的节点我们也可以截取呢，那我们可以看到它还把 relu 的 forwad 进行了截取让它重定位到 symbolic_relu 里面去，接着也是通过 helper 去创建了一个 node，那后面的都是通过这种方式来创建的一个自定义的节点
def export_onnx(model, voxels, coors, batch_size, spatial_shape, save_onnx, save_tensor):
global avoid_reuse_container, tensor_map, nodes, initializers, enable_trace
avoid_reuse_container = []
tensor_map = {}
nodes = []
initializers = []
print("Tracing model inference...")
print("> Do inference...")
with torch.no_grad():
register_tensor(voxels)
enable_trace = True
y = model(voxels, coors, batch_size, spatial_shape)[0]
enable_trace = False
if save_tensor is not None:
print("> Do save tensor, The purpose of this operation is to verify the inference result of C++")
print(f" --> Save inference input voxels to {save_tensor}.voxels, voxels.shape = {voxels.shape}")
funcs.save_tensor(voxels, f"{save_tensor}.voxels")
print(f" --> Save inference input coors to {save_tensor}.coors, coors.shape = {coors.shape}")
funcs.save_tensor(coors, f"{save_tensor}.coors")
print(f" --> Save inference output to {save_tensor}.output, output.shape = {y.shape}")
funcs.save_tensor(y, f"{save_tensor}.output")
print(f" --> Save spatial_shape is {spatial_shape}, batch size is {batch_size}")
print(f" --> Save spatial_shape and batch size to {save_tensor}.info")
funcs.save_tensor([batch_size] + spatial_shape, f"{save_tensor}.info")
print("Tracing done!")
inputs = [
helper.make_value_info(
name="0",
type_proto=helper.make_tensor_type_proto(
elem_type=helper.TensorProto.DataType.FLOAT16,
shape=voxels.size()
)
)
]
outputs = [
helper.make_value_info(
name=get_tensor_id(y),
type_proto=helper.make_tensor_type_proto(
elem_type=helper.TensorProto.DataType.FLOAT16,
shape=y.size()
)
)
]
graph = helper.make_graph(
name="scn",
inputs=inputs,
outputs=outputs,
nodes=nodes,
initializer=initializers
)
opset = [
helper.make_operatorsetid("ai.onnx", 11)
]
model = helper.make_model(graph, opset_imports=opset, producer_name="pytorch", producer_version="1.9")
onnx.save_model(model, save_onnx)
print(f"🚀 The export is completed. ONNX save as {save_onnx} 🤗, Have a nice day~")
# clean memory
avoid_reuse_container = []
tensor_map = {}
nodes = []
initializers = []
我们可以看到最后这些截取到的节点都给放到了 nodes 列表中，在 export_onnx 函数中我们可以看到它通过 helper 去创建了一个 input 一个 output 还创建了一个 graph，我们通过 graph 之后把输入输出以及整个网络中所涉及到的所有节点放在一起形成一个 graph，之后我们再通过 graph 创建一个 model，那最终通过 onnx 的 save_model 就可以把这个 onnx 模型给保存下来了
OK，以上就是带有 spconv 的 SCN 网络导出的一个方式，这里面的重点主要是 register_node 装饰器以及 make_hook 钩子函数，这种语法大家可能平时看到的并不是很多，那这个需要大家自己多看多理解了，另外关于 Python 装饰器的内容大家感兴趣的可以看看：AutoCV第四课：Python基础
这个就是 trace spconv 的一个方式，分析完代码之后我们现在想要导出对应 SCN 的 onnx 该怎么办呢，我们来看下 NVIDIA 官方提供的步骤，如下所示：
1. 从 https://github.com/tianweiy/CenterPoint 下载并配置 CenterPoint 的环境
2. 使用如下指令导出 SCN ONNX
$ cp -r tool/centerpoint-export path/to/CenterPoint
$ cd path/to/CenterPoint
$ python centerpoint-export/export-scn.py --ckpt=epoch_20.pth --save-onnx=scn.nuscenes.onnx
$ cp scn.nuscenes.onnx path/to/3DSparseConvolution/workspace/
那我们首先需要去搭建 CenterPoint 的一个环境，可以按照它的 INSTALL 文件去进行配置，如下图所示：
值得注意的是，CenterPoint 官方提供的环境配置其实是用来做训练用的，所以它这个环境搭建起来会比较大，比较复杂，也会有很多没必要的东西。对于我们来说，我们其实不需要做训练，我们只需要导出 ONNX 就好了，所以我们只要把跟 ONNX 导出有关的包安装下就好了
另外我们可以看到 CenterPoint 官方推荐安装的软件版本都太低了，那 pytorch 版本是 1.1.0 的，cudatoolkit 版本是 10.0 的，现在 cuda 都是 12.0 的，那这个看起来就不太好，那怎么办呢，那这里大家不要按照官方的步骤去配置这个环境，我们自己来配置，那下面是韩老师配置的一个能够导出 ONNX 的最小环境配置，指令如下：
git clone https://github.com/tianweiy/CenterPoint.git
cd CenterPoint
conda create --name centerpoint python=3.7
conda activate centerpoint
conda install pytorch==1.10.1 torchvision==0.11.2 cudatoolkit=11.3 -c pytorch
cd det3d/ops/iou3d_nms
python setup.py build_ext --inplace
pip install numba spconv-cu113 terminaltables addict pyyaml pycocotools onnx
那我们大家按照韩老师提供的这个方式去配置一下环境就好了，大家不要去做 pip install -r requirement.txt，这里面有太多东西没有必要安装，比如说 opencv、matplotlib 等等，每一个都有好几个 G，我们没有必要去安装它，所以就直接跳过就好了
OK，我们按照上述方式将环境配置好之后，再看 NVIDIA 官方提供的导出指令，将 centerpoint-export 文件夹拷贝到 CenterPoint 中，然后执行相关 python 指令导出即可，整个过程如下图所示：
我们可以看到 SCN 的 ONNX 成功导出了，那它效果如下图所示：
从日志信息可以看出它把 SCN 里面的所有节点全都做了一个重定位，导出完之后的 onnx 我们可以使用 Netron 可视化工具打开，如下图所示：
上面就是刚刚导出的 ONNX，它的架构就是这个样子，SparseConvolution 就是我们说的 spconv 节点，它包含的数据信息有很多，那推理的时候我们就只要把这些数据传递给 spconv 推理引擎的 builder 就可以创建一个自定义的节点
那这里给大家讲解了导出 spconv 的代码，讲解了里面的 register_node 和 hook 函数是怎么实现的，这里面大家不清楚的可以去看看视频
OK，那导出带有 spconv 的 SCN 网络的 ONNX 这个部分到这里就结束了，建议大家自己动手把这个 onnx 导出来，同时也可以看一看它的 C++ 推理框架是怎么写的
3. 补充-装饰器+钩子函数
博主这里结合 ChatGPT 的分析把 spconv.forward 如何实现重定位的流程再简单过一遍
我们要分析的代码如下：
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, f"spconv{ilayer}.weight"),
]
if self.bias is not None:
inputs.append(append_initializer(self.bias.data, f"spconv{ilayer}.bias"))
act_type_name = {
tv.gemm.Activation.ReLU : "ReLU",
tv.gemm.Activation.None_ : "None",
tv.gemm.Activation.Sigmoid : "Sigmoid",
tv.gemm.Activation.LeakyReLU : "LeakyReLU"
}
algo_name = {
ConvAlgo.MaskImplicitGemm : "MaskImplicitGemm",
ConvAlgo.MaskSplitImplicitGemm : "MaskSplitImplicitGemm",
ConvAlgo.Native : "Native",
}
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 = act_type_name[self.act_type],
input_shape = x.features.shape,
output_shape = y.features.shape
)
)
首先我们需要知道 @register_node 是一个装饰器，装饰器在 Python 中是一个非常强大的功能，它允许你在不修改原始代码的情况下，给函数增加额外的功能。装饰器本质上是一个函数，它接受一个函数作为参数并返回一个新的函数。当看到 @ 符号用在一个函数定义之前时，它实际上是在应用一个装饰器。
下面我们详细分析一下在执行 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 函数来进行追踪和收集信息。这种机制允许在不更改原始函数代码的前提下，增加额外的功能，这在许多情况下非常有用，特别是在需要追踪或记录函数行为的场景中。
总结
本次课程我们学习了导出带有 spconv 的 SCN 网络的 onnx，首先我们先去看了 spconv 的 C++ 推理框架，由于核心部分未开源我们也只是知道它的基本流程，它主要是通过 protobuf 的函数解析读取的 onnx，然后将解析出来的数据通过 builder 的 push_sparse_conv 函数创建 spconv 的节点，依此类推，SCN 中所有节点均可采用这种方式，这样整个 engine 就构建好了。
接着我们学习了导出 spconv 的 onnx，主要是利用钩子函数对 spconv.forward 进行重定位，重定位到我们自定义的函数中，通过 helper.make_node 创建自定义的节点，将推理所需要的信息全部都填充进去。这样把 SCN 中所有自定义的节点放到 graph 中创建一个 model 保存下来即可
OK，以上就是第 3 小节有关导出带有 spconv 的 SCN 网络的 onnx的全部内容了，下节我们将去学习 spconv 的优化方案—Explicit GEMM Conv，敬请期待😄
下载链接
- 论文下载链接【提取码：6463】
- 数据集下载链接【提取码：data】
- 代码和安装包下载链接【提取码：cuda】
参考
- AutoCV第四课：Python基础
- 复杂onnx解决方案（以sparseconv为例）
- https://github.com/tianweiy/CenterPoint
- https://github.com/NVIDIA-AI-IOT/Lidar_AI_Solution/tree/master/libraries/3DSparseConvolution
