以下文章来源于PandaCV ，作者BBuf
这篇文章从多个角度探索了ONNX，从ONNX的导出到ONNX和Caffe的对比，以及使用ONNX遭遇的困难以及一些解决办法，另外还介绍了ONNXRuntime以及如何基于ONNXRuntime来调试ONNX模型等，后续也会继续结合ONNX做一些探索性工作。
0x0. 前言
这一节我将主要从盘点ONNX模型部署有哪些常见问题，以及针对这些问题提出一些解决方法，另外本文也会简单介绍一个可以快速用于ONNX模型推理验证的框架ONNXRuntime。如果你想用ONNX作为模型转换和部署的工具，可以耐心看下去。今天要讲到的ONNX模型部署碰到的问题大多来自于一些关于ONNX模型部署的文章以及自己使用ONNX进行模型部署过程中的一些经历，有一定的实践意义。
0x1. 导出ONNX
这里以Pytorch为例，来介绍一下要把Pytorch模型导出为ONNX模型需要注意的一些点。首先，Pytorch导出ONNX的代码一般是这样：
发布文章import torch
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = torch.load("test.pth") # pytorch模型加载
batch_size = 1 #批处理大小
input_shape = (3, 244, 224) #输入数据,改成自己的输入shape
# #set the model to inference mode
model.eval()
x = torch.randn(batch_size, *input_shape) # 生成张量
x = x.to(device)
export_onnx_file = "test.onnx" # 目的ONNX文件名
torch.onnx.export(model
x,
export_onnx_file,
opset_version=10,
do_constant_folding=True, # 是否执行常量折叠优化
input_names=["input"], # 输入名
output_names=["output"], # 输出名
dynamic_axes={"input":{0:"batch_size"}, # 批处理变量
"output":{0:"batch_size"}})可以看到Pytorch提供了一个ONNX模型导出的专用接口，只需要配置好相关的模型和参数就可以完成自动导出ONNX模型的操作了。代码相关细节请自行查看，这里来列举几个导出ONNX模型中应该注意的问题。
自定义OP问题
以2020年的YOLOV5为例，在模型的BackBone部分自定义了一个Focus OP，这个OP的代码实现为：
class Focus(nn.Module): # Focus wh information into c-space def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True): # ch_in, ch_out, kernel, stride, padding, groups super(Focus, self).__init__() self.conv = Conv(c1 * 4, c2, k, s, p, g, act) # self.contract = Contract(gain=2) def forward(self, x): # x(b,c,w,h) -> y(b,4c,w/2,h/2) return self.conv(torch.cat([x[..., ::2, ::2], x[..., 1::2, ::2], x[..., ::2, 1::2], x[..., 1::2, 1::2]], 1)) # return self.conv(self.contract(x))
这个操作就是一个stride slice然后再concat的操作，类似于PixelShuffle的逆向过程。下面是把YOLOv5模型导出ONNX模型之后Focus层的可视化结果。
可以看到这个OP在使用Pytorch导出ONNX的过程中被拆成了很多更小的操作，这个时候Focus OP的问题就是推理的效率可能比较低并且拆成的小OP各个推理框架的支持程度不一致。要解决这种问题，要么直接在前向推理框架实现一个自定义的Focus OP，ncnn在实现yolov5的时候也是这样做的：
https://github.com/Tencent/ncnn/blob/master/examples/yolov5.cpp#L24。
要么将这个OP使用其它的操作来近似代替，比如这里可以使用一个stride为2的卷积OP来代替Focus结构，注意代替之后有可能准确率会下降，需要做精度和部署友好性的平衡。
综上，自定义的OP在导出ONNX进行部署时，除了考虑ONNX模型的执行效率问题，还要考虑框架是否支持的问题。如果想快速迭代产品，建议尽量以一些经典结构为基础，尽量少引入自定义OP。
后处理问题
如果我们要导出检测网络的ONNX模型进行部署，就会碰到这个问题，后处理部分是否需要导入到ONNX模型？
我们知道在使用Pytorch导出ONNX模型时，所有的Aten操作都会被ONNX记录下来（具体记录什么内容请参考文章开头链接推文的介绍），成为一个DAG。然后ONNX会根据这个DAG的输出节点来反推这个DAG中有哪些节点是有用的，这样获得的就是最终的ONNX模型。而后处理，比如非极大值抑制也是通过Aten操作拼起来的，所谓Aten操作就是Pytorch中的基础算术单元比如加减乘除，所有的OP以及和Tensor相关的操作都基于Aten中的操作拼。
检测网络比如YOLOV3的后处理就是NMS，代码示例如
https://github.com/ultralytics/yolov3/blob/master/utils/general.py#L325。
当我们完成检测网络的训练之后直接导出ONNX模型我们就会发现NMS这个实现也全部被导入了ONNX，如下图所示：
这个结构非常复杂，我们要在实际业务中去部署这个模型难度是很大的。另外，刚才我们提到ONNX模型只能记录Pytorch中的Aten操作，对其它的一些逻辑运算符比如if是无能为力的（意思是不能记录if的多个子图），而后处理过程中根据置信度阈值来筛选目标框是常规操作。如果我们在导出ONNX模型时是随机输入或者没有指定目标的图片就会导致这个ONNX记录下来的DAG可能有缺失。最后，每个人实现后处理的方式可能都是不一样的，这也增加了ONNX模型部署的难度。为了部署的友好性和降低转换过程中的风险，后处理过程最好由读者自己完成，我们只需要导出模型的Backbone和Neck部分为ONNX。
具体来说，我们只需要在Pytorch的代码实现中屏蔽掉后处理部分然后导出ONNX模型即可。这也是目前使用ONNX部署检测模型的通用方案。
所以，针对后处理问题，我们的结论就是在使用ONNX进行部署时直接屏蔽后处理，将后处理单独拿出来处理。
胶水OP问题。
在导出ONNX模型的过程中，经常会带来一些胶水OP，比如Gather, Shape等等。例如上节推文中介绍到当执行下面的Pytorch导出ONNX程序时，就会引入很多胶水OP。
import torch class JustReshape(torch.nn.Module): def __init__(self): super(JustReshape, self).__init__() def forward(self, x): return x.view((x.shape[0], x.shape[1], x.shape[3], x.shape[2])) net = JustReshape() model_name = '../model/just_reshape.onnx' dummy_input = torch.randn(2, 3, 4, 5) torch.onnx.export(net, dummy_input, model_name, input_names=['input'], output_names=['output'])
导出的ONNX模型可视化如下：
这个时候的做法一般就是过一遍onnx-simplifer，可以去除这些胶水OP获得一个简化后的模型。
综上，我们在导出ONNX模型的一般流程就是，去掉后处理，尽量不引入自定义OP，然后导出ONNX模型，并过一遍大老师的https://github.com/daquexian/onnx-simplifier，这样就可以获得一个精简的易于部署的ONNX模型。从ONNX官方仓库提供的模型来看，似乎微软真的想用ONNX来统一所有框架的所有操作。但理想很丰满，现实很骨干，各种训练框架的数据排布，OP实现不一致，人为后处理不一致，各种推理框架支持度不一致，推理芯片SDK的OP支持度不一致都让这个ONNX（万能格式）遭遇了困难，所以在基于ONNX做一些部署业务的时候，也要有清晰的判断并选取风险最小的方法。
0x2. ONNX or Caffe？
这个问题其实源于之前做模型转换和基于TensorRT部署一些模型时候的思考。我们还是以Pytorch为例，要把Pytorch模型通过TensorRT部署到GPU上，一般就是Pytorch->Caffe->TensorRT以及Pytorch->ONNX->TensorRT（当然Pytorch也是支持直接转换到TensorRT，这里不关心）。那么这里就有一个问题，我们选择哪一条路比较好？
其实，我想说的应该是Caffe是过去，而ONNX是将来。为什么要这样说？
首先很多国产推理芯片比如海思NNIE，高通SNPE它们首先支持的都是Caffe这种模型格式，这可能是因为年代的原因，也有可能是因为这些推理SDK实现的时候OP都非常粗粒度。比如它对卷积做定制的优化，有NC4HW4，有Im2Col+gemm，有Winograd等等非常多方法，后面还考虑到量化，半精度等等，然后通过给它喂Caffe模型它就知道要对这个网络里面对应的卷积层进行硬件加速了。所以这些芯片支持的网络是有限的，比如我们要在Hisi35xx中部署一个含有upsample层的Pytorch模型是比较麻烦的，可能不太聪明的工程说我们要把这个模型回退给训练人员改成支持的上采样方式进行训练，而聪明的工程师可能说直接把upsample的参数填到反卷积层的参数就可以了。无论是哪种方式都是比较麻烦的，所以Caffe的缺点就是灵活度太差。其实基于Caffe进行部署的方式仍然在工业界发力，ONNX是趋势，但是ONNX现在还没有完全取代Caffe。
接下来，我们要再提一下上面那个if的事情了，假设现在有一个新的SOTA模型被提出，这个模型有一个自定义的OP，作者是用Pytorch的Aten操作拼的，逻辑大概是这样：
result = check() if result == 0: result = algorithm1(result) else: result = algorithm2(result) return result
然后考虑将这个模型导出ONNX或者转换为Caffe，如果是Caffe的话我们需要去实现这个自定义的OP，并将其注册到Caffe的OP管理文件中，虽然这里比较繁琐，但是我们可以将if操作隐藏在这个大的OP内部，这个if操作可以保留下来。而如果我们通过导出ONNX模型的方式if子图只能保留一部分，要么保留algorithm1，要么保留algorithm2对应的子图，这种情况ONNX似乎就没办法处理了。这个时候要么保存两个ONNX模型，要么修改算法逻辑绕过这个问题。从这里引申一下，如果我们碰到有递归关系的网络，基于ONNX应当怎么部署？ONNX还有一个缺点就是OP的细粒度太细，执行效率低，不过ONNX已经推出了多种化方法可以将OP的细粒度变粗，提高模型执行效率。目前在众多经典算法上，ONNX已经支持得非常好了。
最后，越来越多的厂商推出的端侧推理芯片开始支持ONNX，比如地平线的BPU，华为的Ascend310相关的工具链都开始接入ONNX，所以个人认为ONNX是推理框架模型转换的未来，不过仍需时间考验，毕竟谁也不希望因为框架OP对齐的原因导出一个超级复杂的ONNX模型，还是简化不了的那种，导致部署难度很大。
0x3. 一些典型坑点及解决办法
第一节已经提到，将我们的ONNX模型过一遍onnx-simplifer之后就可以去掉胶水OP并将一些细粒度的OP进行op fuse成粗粒度的OP，并解决一部分由于Pytorch和ONNX OP实现方式不一致而导致模型变复杂的问题。除了这些问题，本节再列举一些ONNX模型部署中容易碰到的坑点，并尝试给出一些解决办法。
预处理问题。
和后处理对应的还有预处理问题，如果在Pytorch中使用下面的代码导出ONNX模型。
import torch class JustReshape(torch.nn.Module): def __init__(self): super(JustReshape, self).__init__() self.mean = torch.randn(2, 3, 4, 5) self.std = torch.randn(2, 3, 4, 5) def forward(self, x): x = (x - self.mean) / self.std return x.view((x.shape[0], x.shape[1], x.shape[3], x.shape[2])) net = JustReshape() model_name = '../model/just_reshape.onnx' dummy_input = torch.randn(2, 3, 4, 5) torch.onnx.export(net, dummy_input, model_name, input_names=['input'], output_names=['output'])
我们先给这个ONNX模型过一遍onnx-simplifer，然后使用Netron可视化之后模型大概长这样：
如果我们要把这个模型放到NPU上部署，如果NPU芯片不支持Sub和Div的量化计算，那么这两个操作会被回退到NPU上进行计算，这显然是不合理的，因为我们总是想网络在NPU上能一镜到底，中间断开必定会影响模型效率，所以这里的解决办法就是把预处理放在基于nn.Module搭建模型的代码之外，然后推理的时候先把预处理做掉即可。
*博客内容为网友个人发布，仅代表博主个人观点，如有侵权请联系工作人员删除。
