05 感知 3D目标检测 模型导出和部署（基础知识）

1. 模型导出基础介绍

• torch.onnx.export中需要的模型实际上是一个torch.jit.ScriptModule。而要把普通 PyTorch 模型转一个这样的 TorchScript 模型，有跟踪（trace）和记录（script）两种导出计算图的方法。如果给torch.onnx.export传入了一个普通 PyTorch 模型（torch.nn.Module)，那么这个模型会默认使用跟踪的方法导出。
• 跟踪法只能通过实际运行一遍模型的方法导出模型的静态图，即无法识别出模型中的控制流（如循环）；记录法则能通过解析模型来正确记录所有的控制流。

2. 是否有些pytorch算子在onnx中不支持？如何在pytorch支持更多onnx算子？

Ref 1. 模型部署入门教程（四）：在 PyTorch 中支持更多 ONNX 算子: https://zhuanlan.zhihu.com/p/513387413

1. 在最简单的情况下，我们只要把 PyTorch 算子的输入用g.op()一一对应到 ONNX 算子上即可，并把g.op()的返回值作为符号函数的返回值。在情况更复杂时，我们转换一个 PyTorch 算子可能要新建若干个 ONNX 算子。

• 为算子添加符号函数一般要经过以下几步：1. 获取原算子的前向推理接口；2. 获取目标 ONNX 算子的定义；3. 编写符号函数并绑定。

2. torch.autograd.Function 能完成算子实现和算子调用的隔离。不管算子是怎么实现的，它封装后的使用体验以及 ONNX 导出方法会和原生的 PyTorch 算子一样。这是我们比较推荐的为算子添加 ONNX 支持的方法。

1. 为torch添加c++扩展

"""
这里的 my_lib 是我们未来要在 Python 里导入的模块名。双引号中的 my_add 是 Python 调用接口的名称
"""
// my_add.cpp 
 
#include <torch/torch.h> 
 
torch::Tensor my_add(torch::Tensor a, torch::Tensor b) 
{ 
    return 2 * a + b; 
} 
 
PYBIND11_MODULE(my_lib, m) 
{ 
    m.def("my_add", my_add); 
} 

2. 编写如下的 Python 代码并命名为 "setup.py"，来编译刚刚的 C++ 文件

from setuptools import setup 
from torch.utils import cpp_extension 
 
setup(name='my_add', 
      ext_modules=[cpp_extension.CppExtension('my_lib', ['my_add.cpp'])], 
      cmdclass={'build_ext': cpp_extension.BuildExtension}) 

3. 直接用 Python 接口调用 C++ 函数不太“美观”，一种比较优雅的做法是把这个调用接口封装起来。这里我们用 torch.autograd.Function 来封装算子的底层调用

import torch 
import my_lib 
class MyAddFunction(torch.autograd.Function): 
 
    @staticmethod 
    def forward(ctx, a, b): 
        return my_lib.my_add(a, b) 
 
    @staticmethod 
    def symbolic(g, a, b): 
        two = g.op("Constant", value_t=torch.tensor([2])) 
        a = g.op('Mul', a, two) 
        return g.op('Add', a, b) 

• 在 forward 函数中，我们用 my_lib.my_add(a, b) 就可以调用之前写的C++函数了。这里 my_lib 是库名，my_add 是函数名，这两个名字是在前面C++的 PYBIND11_MODULE 中定义的。
• 在 ONNX 中，我们需要把新建常量当成一个算子来看待，尽管这个算子并不会以节点的形式出现在 ONNX 模型的可视化结果里。

---

补充：
1. 符号函数 symbolic ？

• 符号函数，可以看成是 PyTorch 算子类的一个静态方法。在把 PyTorch 模型转换成 ONNX 模型时，各个 PyTorch 算子的符号函数会被依次调用，以完成 PyTorch 算子到 ONNX 算子的转换。第一个参数就固定叫 g，它表示和计算图相关的内容；后面的每个参数都表示算子的输入，需要和算子的前向推理接口的输入相同。
• g 有一个方法 op。在把 PyTorch 算子转换成 ONNX 算子时，需要在符号函数中调用此方法来为最终的计算图添加一个 ONNX 算子。其中，第一个参数是算子名称。如果该算子是普通的 ONNX 算子，只需要把它在 ONNX 官方文档里的名称填进去即可。

def symbolic(g: torch._C.Graph, input_0: torch._C.Value, input_1: torch._C.Value, ...): 
	pass 
	
def op(name: str, input_0: torch._C.Value, input_1: torch._C.Value, ...) :
	pass
	

2. torch.autograd.Function介绍？

• Function 类本身表示 PyTorch 的一个可导函数，只要为其定义了前向推理和反向传播的实现，我们就可以把它当成一个普通 PyTorch 函数来使用。
• PyTorch 会自动调度该函数，合适地执行前向和反向计算。对模型部署来说，Function 类有一个很好的性质：如果它定义了 symbolic 静态方法，该 Function 在执行 torch.onnx.export() 时就可以根据 symbolic 中定义的规则转换成 ONNX 算子。这个 symbolic 就是前面提到的符号函数，只是它的名称必须是 symbolic 而已。
• apply是torch.autograd.Function 的一个方法，这个方法完成了 Function 在前向推理或者反向传播时的调度。我们在使用 Function 的派生类做推理时，不应该显式地调用 forward()，而应该调用其 apply 方法。

3. torch.autograd.Function backward函数如何定义？

3. 关于onnx，pytorch和onnx精度对齐

Ref 1. 模型部署入门教程（五）：ONNX 模型的修改与调试: https://zhuanlan.zhihu.com/p/516920606
Ref 2. 模型部署入门教程（六）：实现 PyTorch-ONNX 精度对齐工具: https://zhuanlan.zhihu.com/p/543973749 待细读

ONNX 在底层是用 Protobuf 定义的。Protobuf，全称 Protocol Buffer，是 Google 提出的一套表示和序列化数据的机制。使用 Protobuf 时，用户需要先写一份数据定义文件，再根据这份定义文件把数据存储进一份二进制文件。可以说，数据定义文件就是数据类，二进制文件就是数据类的实例。

4. onnx转TensorRT，遇到不能转换的算子，如何自定义实现转换？

在具体的生产环境中，ONNX 模型常常需要被转换成能被具体推理后端使用的模型格式。

补充：

1. CUDA vs CuDNN vs TensorRT，英伟达这些库的关系？
在训练神经网络时，框架如TensorFlow或PyTorch会调用cuDNN中的函数，而这些函数又依赖于CUDA的底层接口。所以cuDNN实际上是建立在CUDA之上的，专门为深度学习优化的库，两者协同工作来加速计算。CUDA是基础，cuDNN加速训练中的特定操作，TensorRT优化推理阶段的性能。