torch.onnx.export 详解
1. 计算图导出方法
TorchScript 是一种序列化和优化 PyTorch 模型的格式，在优化过程中，一个torch.nn.Module 模型会被转换成 TorchScript 的 torch.jit.ScriptModule 模型。现在，TorchScript 也被常当成一种中间表示使用。
torch.onnx.export 中需要的模型实际上是一个 torch.jit.ScriptModule。而要把普通 PyTorch 模型转一个这样的 TorchScript 模型，有跟踪（trace）和记录（script）两种导出计算图的方法。如果给 torch.onnx.export 传入了一个普通 PyTorch 模型（torch.nn.Module)，那么这个模型会默认使用跟踪的方法导出。
跟踪法只能通过实际运行一遍模型的方法导出模型的静态图，即无法识别出模型中的控制流（如循环）；记录法则能通过解析模型来正确记录所有的控制流。
推理引擎对静态图的支持更好，通常我们在模型部署时不需要显式地把 PyTorch 模型转成 TorchScript 模型，直接把 PyTorch 模型用 torch.onnx.export 跟踪导出即可。在模型转换报错时需要注意是否发生在 PyTorch 转 TorchScript 阶段。
2.参数讲解
def export(model, # 模型
args, # 模型的输入
f, # 导出ONNX的名字
export_params=True, # 模型中是否存储模型权重。ONNX 是用同一个文件表示记录模型的结构和权重的。
verbose=False,
training=TrainingMode.EVAL,
input_names=None, # 输入张量
output_names=None, # 输出张量
aten=False,
export_raw_ir=False,
operator_export_type=None,
opset_version=None, # 转换时参考哪个 ONNX 算子集版本
_retain_param_name=True,
do_constant_folding=True,
example_outputs=None,
strip_doc_string=True,
dynamic_axes=None, # 指定输入输出张量的哪些维度是动态的。
keep_initializers_as_inputs=None,
custom_opsets=None,
enable_onnx_checker=True,
use_external_data_format=False)
dynamic_axes 的设置例子
dynamic_axes_0 = {
'in' : [0],
'out' : [0]
}
# ONNX 要求每个动态维度都有一个名字，这样写的话会引出一条 UserWarning，警告我们通过列表的方式设置动态维度的话系统会自动为它们分配名字。
dynamic_axes_0 = {
'in' : {0: 'batch'},
'out' : {0: 'batch'}
}
torch.onnx.export(model,
dummy_input, model_names[1],
input_names=['in'],
output_names=['out'],
dynamic_axes=dynamic_axes_0)
3.使用技巧
(1) 使模型在 ONNX 转换时有不同的行为
有些时候，我们希望模型在直接用 PyTorch 推理时有一套逻辑，而在导出的 ONNX 模型中有另一套逻辑。比如，我们可以把一些后处理的逻辑放在模型里，以简化除运行模型之外的其他代码。torch.onnx.is_in_onnx_export() 可以实现这一任务，该函数仅在执行 torch.onnx.export() 时为真。
(2) 利用中断张量跟踪的操作
PyTorch 转 ONNX 的跟踪导出法会把某些取决于输入的中间结果变成常量，从而使导出的 ONNX 模型和原来的模型有出入。涉及张量与普通变量转换的逻辑都会导致最终的 ONNX 模型不太正确。
4.如何判断某个 PyTorch 算子在 ONNX 中是否兼容
在转换普通的 torch.nn.Module 模型时，PyTorch 一方面会用跟踪法执行前向推理，把遇到的算子整合成计算图；另一方面，PyTorch 还会把遇到的每个算子翻译成 ONNX 中定义的算子。在这个翻译过程中，可能会碰到以下情况：
(1) 该算子可以一对一地翻译成一个 ONNX 算子。
(2) 该算子在 ONNX 中没有直接对应的算子，会翻译成一至多个 ONNX 算子。
(3) 该算子没有定义翻译成 ONNX 的规则，报错。
PyTorch 对 ONNX 算子的映射
torch.onnx 目录网址：https://github.com/pytorch/pytorch/tree/master/torch/onnx
其中，symbolic_opset{n}.py（符号表文件）即表示 PyTorch 在支持第 n 版 ONNX 算子集时新加入的内容。以 “bicubic” 为例，使用搜索功能，在 torch/onnx 文件夹搜索 “bicubic”，可以发现这个这个插值在第 11 个版本的定义文件中
实战
PyTorch 算子顺利转换到 ONNX需要保证以下三个环节都不出错：
算子在 PyTorch 中有实现
有把该 PyTorch 算子映射成一个或多个 ONNX 算子的方法
ONNX 有相应的算子
1.支持 ATen 算子
算子在 ATen 中已经实现了，ONNX 中也有相关算子的定义，但是相关算子映射成 ONNX 的规则没有写。在这种情况下，我们只需要为 ATen 算子补充描述映射规则的符号函数就行了。
ATen: PyTorch 内置的 C++ 张量计算库，PyTorch 算子在底层绝大多数计算都是用 ATen 实现的
Asinh 算子在 ATen 中有实现，却缺少了映射到 ONNX 算子的符号函数。我们来尝试为它补充符号函数，并导出一个包含这个算子的 ONNX 模型。
获取 ATen 中算子接口定义
1. 为了编写符号函数，我们需要获得 asinh 推理接口的输入参数定义。这时，我们要去 torch/_C/_VariableFunctions.pyi 和 torch/nn/functional.pyi 这两个文件中搜索我们刚刚得到的这个算子名。
def atanh(input: Tensor, *, out: Optional[Tensor]=None) -> Tensor: ...
import torch
class Model(torch.nn.Module):
def __init__(self):
super().__init__()
def forward(self, x):
return torch.asinh(x)
运行上面的代码不会报错，算子在 ATen 中已经实现了。
model = Model()
input = torch.rand(1, 3, 10, 10)
torch.onnx.export(model, input, 'asinh.onnx', opset_version=9)
进行转换的时候会报错：UnsupportedOperatorError: Exporting the operator ::asinh to ONNX opset version 9 is not supported. Please feel free to request support or submit a pull request on PyTorch GitHub.
Asinh 算子出现于第 9 个 ONNX 算子集。 没有相应的映射规则所以报错
编写符号函数(映射规则)
符号函数，可以看成是 PyTorch 算子类的一个静态方法。在把 PyTorch 模型转换成 ONNX 模型时，各个 PyTorch 算子的符号函数会被依次调用，以完成 PyTorch 算子到 ONNX 算子的转换。
def symbolic(g: torch._C.Graph, input_0: torch._C.Value, input_1: torch._C.Value, ...):
第一个参数就固定叫 g，它表示和计算图相关的内容；后面的每个参数都表示算子的输入，需要和算子的前向推理接口的输入相同。对于 ATen 算子来说，它们的前向推理接口就是上述两个 .pyi 文件里的函数接口。
g 有一个方法 op。在把 PyTorch 算子转换成 ONNX 算子时，需要在符号函数中调用此方法来为最终的计算图添加一个 ONNX 算子。
def op(name: str, input_0: torch._C.Value, input_1: torch._C.Value, ...)
第一个参数是ONNX算子名称。Asinh 是普通的 ONNX 算子，只需要把它在 ONNX 官方文档里的名称填进去即可 Asinh 的文档写道：该算子有一个输入 input，一个输出 output，二者的类型都为张量。
在上一小节得知了 asinh 的推理接口定义，在这一小节里收集了 ONNX 算子 Asinh 的定义。现在，我们可以用代码来补充这二者的映射关系了。
from torch.onnx.symbolic_registry import register_op
# asinh_symbolic就是asinh的符号函数。从除g以外的第二个输入参数开始，其输入参数应该严格对应它在 ATen 中的定义
def asinh_symbolic(g, input, *, out=None):
return g.op("Asinh", input)
# 把这个符号函数和原来的 ATen 算子“绑定”起来
# register_op的第一个参数是目标 ATen 算子名，
# 第二个是要注册的符号函数，这两个参数很好理解。
# 第三个参数是算子的“域”，对于普通 ONNX 算子，直接填空字符串即可。
# 第四个参数表示向哪个算子集版本注册。
register_op('asinh', asinh_symbolic, '', 9)
第一个参数"Asinh"是算子在 ONNX 中的名称。至于第二个参数 input，如我们刚刚在文档里所见，这个算子只有一个输入，因此我们只要把符号函数的输入参数 input 对应过去就行。ONNX 的 Asinh 的输出和 ATen 的 asinh 的输出是一致的，因此我们直接把 g.op() 的结果返回即可。
model = Model()
input = torch.rand(1, 3, 10, 10)
torch.onnx.export(model, input, 'asinh.onnx', opset_version=9)
测试算子
import onnxruntime
import torch
import numpy as np
class Model(torch.nn.Module):
def __init__(self):
super().__init__()
def forward(self, x):
return torch.asinh(x)
model = Model()
input = torch.rand(1, 3, 10, 10)
torch_output = model(input).detach().numpy()
sess = onnxruntime.InferenceSession('asinh.onnx')
ort_output = sess.run(None, {'onnx::Asinh_0': input.numpy()})[0]
# 使用 np.allclose 来保证两个结果张量的误差在一个可以允许的范围内
assert np.allclose(torch_output, ort_output)
2.支持 TorchScript 算子
对于一些比较复杂的运算，仅使用 PyTorch 原生算子是无法实现的。就需要自定义一个 PyTorch 算子，再把它转换到 ONNX 中。PyTorch 官方比较推荐的一种做法是添加 TorchScript 算子。
为算子添加符号函数一般要经过以下几步：
1）获取原算子的前向推理接口。
2）获取目标 ONNX 算子的定义。
3）编写符号函数并绑定。
使用 TorchScript 算子
新增 TorchScript 算子比较复杂，就以可变形卷积（Deformable Convolution）算子为例，为其添加 ONNX 支持
import torch
import torchvision
class Model(torch.nn.Module):
def __init__(self):
super().__init__()
self.conv1 = torch.nn.Conv2d(3, 18, 3)
self.conv2 = torchvision.ops.DeformConv2d(3, 3, 3)
def forward(self, x):
return self.conv2(x, self.conv1(x))
查询算子的前向推理接口。DeformConv2d 层最终会调用 deform_conv2d 这个算子。可以在 torchvision/csrc/ops/deform_conv2d.py 中查到该算子的调用接口：
def deform_conv2d(
input: Tensor,
offset: Tensor,
weight: Tensor,
bias: Optional[Tensor] = None,
stride: Tuple[int, int] = (1, 1),
padding: Tuple[int, int] = (0, 0),
dilation: Tuple[int, int] = (1, 1),
mask: Optional[Tensor] = None,
) -> Tensor:
自定义 ONNX 算子
g.op() 是用来定义 ONNX 算子的函数。对于 ONNX 官方定义的算子，g.op() 的第一个参数就是该算子的名称。而对于一个自定义算子，g.op() 的第一个参数是一个带命名空间的算子名，比如：
g.op(“custom::deform_conv2d”, …)
其中，"::"前面的内容就是我们的命名空间。该概念和 C++ 的命名空间类似，是为了防止命名冲突而设定的。如果在 g.op() 里不加前面的命名空间，则算子会被默认成 ONNX 的官方算子。
PyTorch 在运行 g.op() 时会对官方的算子做检查，如果算子名有误，或者算子的输入类型不正确， g.op() 就会报错。为了让我们随心所欲地定义新 ONNX 算子，我们必须设定一个命名空间，给算子取个名，再定义自己的算子。
在后续的文章中，我们再介绍在各个推理引擎中添加新 ONNX 算子支持的方法。此处，我们只关心如何导出一个包含新 ONNX 算子节点的 onnx 文件。因此，我们可以为新算子编写如下简单的符号函数：
from torch.onnx import register_custom_op_symbolic
from torch.onnx.symbolic_helper import parse_args
# TorchScript 算子的符号函数要求标注出每一个输入参数的类型。
# 比如"v"表示 Torch 库里的 value 类型，一般用于标注张量，而"i"表示 int 类型，"f"表示 float 类型，"none"表示该参数为空。
# 具体的类型含义可以在 torch.onnx.symbolic_helper.py
@parse_args("v", "v", "v", "v", "v", "i", "i", "i", "i", "i", "i", "i", "i", "none")
def symbolic(g,
input,
weight,
offset,
mask,
bias,
stride_h, stride_w,
pad_h, pad_w,
dil_h, dil_w,
n_weight_grps,
n_offset_grps,
use_mask):
return g.op("custom::deform_conv2d", input, offset,
stride_h_i = stride_h, stride_w_i = stride_w)
# 注册符号函数
# 前面的 register_op 类似，注册符号函数时，我们要输入算子名、符号函数、算子集版本。
# 与前面不同的是，这里的算子集版本是最早生效版本，在这里设定版本 9，意味着之后的版本集都能使用这个新算子。
register_custom_op_symbolic("torchvision::deform_conv2d", symbolic, 9)
model = Model()
input = torch.rand(1, 3, 10, 10)
torch.onnx.export(model, input, 'dcn.onnx')
使用 torch.autograd.Function
为 PyTorch 添加简单的 C++ 拓展还是很方便的。对于我们定义的 my_add 算子，可以用以下的 C++ 源文件来实现。我们把该文件命名为 “my_add.cpp”：
// 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);
}
torch::Tensor 就是 C++ 中 torch 的张量类型，它的加法和乘法等运算符均已重载。因此，我们可以像对普通标量一样对张量做加法和乘法。
用 PYBIND11_MODULE 来为 C++ 函数提供 Python 调用接口。这里的 my_lib 是我们未来要在 Python 里导入的模块名。双引号中的 my_add 是 Python 调用接口的名称，这里我们对齐 C++ 函数的名称，依然用 "my_add"这个名字。
编写如下的 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})
这段代码使用了 Python 的 setuptools 编译功能和 PyTorch 的 C++ 拓展工具函数，可以编译包含了 torch 库的 C++ 源文件。这里我们需要填写的只有模块名和模块中的源文件名。我们刚刚把模块命名为 my_lib，而源文件只有一个 my_add.cpp，因此拓展模块那一行要写成 ext_modules=[cpp_extension.CppExtension(‘my_lib’, [‘my_add.cpp’])]。
执行 python setup.py develop 编译
用 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)
Function 类本身表示 PyTorch 的一个可导函数，只要为其定义了前向推理和反向传播的实现，我们就可以把它当成一个普通 PyTorch 函数来使用。
PyTorch 会自动调度该函数，合适地执行前向和反向计算。对模型部署来说，Function 类有一个很好的性质：如果它定义了 symbolic 静态方法，该 Function 在执行 torch.onnx.export() 时就可以根据 symbolic 中定义的规则转换成 ONNX 算子。这个 symbolic 就是前面提到的符号函数，只是它的名称必须是 symbolic 而已。
在 forward 函数中，我们用 my_lib.my_add(a, b) 就可以调用之前写的C++函数了。这里 my_lib 是库名，my_add 是函数名，这两个名字是在前面C++的 PYBIND11_MODULE 中定义的。
在 ONNX 中，我们需要把新建常量当成一个算子来看待，尽管这个算子并不会以节点的形式出现在 ONNX 模型的可视化结果里。
# apply是torch.autograd.Function 的一个方法，这个方法完成了 Function 在前向推理或者反向传播时的调度。
# 在使用 Function 的派生类做推理时，不应该显式地调用 forward()，而应该调用其 apply 方法。
my_add = MyAddFunction.apply
# MyAdd(torch.nn.Module)封装了my_add，就和封装了conv2d 的 torch.nn.Conv2d 一样。
class MyAdd(torch.nn.Module):
def __init__(self):
super().__init__()
def forward(self, a, b):
# my_add 的地位，和 PyTorch 的 asinh, interpolate, conv2d等原生函数是类似的
return my_add(a, b)
1500
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
