模型部署翻车记：pytorch转onnx踩坑实录

本文原创首发于极市平台公众号，如需转载请私信作者

在深度学习模型部署时，从pytorch转换onnx的过程中，踩了一些坑。本文总结了这些踩坑记录，希望可以帮助其他人。

首先，简单说明一下pytorch转onnx的意义。在pytorch训练出一个深度学习模型后，需要在TensorRT或者openvino部署，这时需要先把Pytorch模型转换到onnx模型之后再做其它转换。因此，在使用pytorch训练深度学习模型完成后，在TensorRT或者openvino或者opencv和onnxruntime部署时，pytorch模型转onnx这一步是必不可少的。接下来通过几个实例程序，介绍pytorch转换onnx的过程中遇到的坑。

1.  opencv里的深度学习模块不支持3维池化层

起初，我在微信公众号里看到一篇文章《使用Python和YOLO检测车牌》。文中展示的检测结果如下，其实这种检测结果并不是一个优良的结果，可以看到检测框里的车牌是倾斜的，如果要识别车牌里的文字，那么倾斜的车牌会严重影响车牌识别结果的。

对于车牌识别这种场景，在做车牌检测时，一种优良的检测结果应该是这样的，如下图所示。

在输出车牌检测框的同时输出检测到的车牌的4个角点。有了这4个角点之后，对车牌做透视变换，这时的车牌就是水平放置的，最后做车牌识别，这样就做成了一个车牌识别系统，在这个系统里包含车牌检测，车牌矫正，车牌识别三个模块。车牌检测模块使用retinaface，原始的retinaface是做人脸检测的，它能输出人脸检测矩形框和人脸5个关键点。考虑到车牌只有4个点，于是修改retinaface的网络结构使其输出4个关键点，然后在车牌数据集训练，训练完成后，以一幅图片上做目标检测的结果如上图所示。车牌矫正模块使用了传统图像处理方法，关键函数是opencv里的getPerspectiveTransform和warpPerspective。车牌识别模块使用Intel公司提出的LPRNet。

整套程序是基于pytorch框架运行的，我把这套程序发布在github上，地址是 https://github.com/hpc203/license-plate-detect-recoginition-pytorch

接下来我就尝试把pytorch模型转换到onnx文件，然后使用opencv做车牌检测与识别。然而在转换完成onnx文件后，使用opencv读取onnx文件遇到了一些坑，我在网上搜索，也没有找到解决办法。

转换过程分两步，首先是转换车牌检测retinaface到onnx文件，这一步倒是很顺利，转换没有出错，并且使用opencv读取onnx文件做前向推理的输出结果也是正确的。第二步转换车牌识别LPRNet到onnx文件，由于Pytorch自带torch.onnx.export转换得到的ONNX，因此转换的代码很简单，在生成onnx文件后，opencv读取onnx文件出现了模型其妙的错误。程序运行的结果截图如下

从打印结果看，torch.onnx.export生成onnx文件时没有问题的，但是在cv2.dnn.readNet这一步出现异常导致程序中断，并且打印出的异常信息是一连串的数字，去百度搜索也么找到解决办法。观察LPRNet的网络结构，发现在LPRNet里定义了3维池化层，代码截图如下

于是，我做了一个实验，定义一个只含有3维池化层的网络，转换生成onnx文件，然后opencv读取onnx文件做前向推理，程序运行结果如下。

可以看到在这时能成功读取onnx文件，但是在执行前向计算model.forward时出错，换成3维平均池化，运行结果如下

可以看到依然出错，这说明opencv的深度学习模块里不支持3维池化。不过，对比3维池化和2维池化的前向计算原理可以发现，3维池化其实等价于2个2维池化。程序实例如下

程序运行结果打印信息是相等。从这里就可以看出opencv里的深度学习模块并不支持3维池化的前向计算，这期待后续新版本的opencv里能添加3维池化的计算。这时在LPRNet网络结构定义文件里修改3维池化层，重新生成onnx文件，opencv读取onnx文件执行前向计算后依然出错，运行结果如下。

于是继续观察LPRNet的网络结构，在forward函数里看到有求平均值的操作，代码截图如下所示

注意到第一个torch.mean函数里没有声明在哪个维度求平均值，这说明它是对一个4维四维张量的整体求平均值，这时候从一个4维空间搜索成一个点，也就是一个标量数值。但是在pytorch里，对一个张量求平均值后依然是一个张量，只不过它的维度shape是空的，示例代码如下。这时如果想要访问平均值，需要加上.item()，这个是需要注意的一个pytorch知识点。

在修改这个代码bug后重新生成onnx文件，使用opencv读取onnx文件做前向计算就不再出现异常错误了。

通过以上几个程序实验，可以总结出opencv读取onnx文件做深度学习前向计算的2个坑：

(1) .opencv里的深度学习模块不支持3维池化计算，解决办法是修改原始网络结构，把3维池化转换成两个2维池化，重新生成onnx文件

(2) .当神经网络里有torch.mean和torch.sum这种把4维张量收缩到一个数值的运算时，opencv执行forward会出错，这时的解决办法是修改原始网络结构，在torch.mean的后面加上.item()

此外，在torch.mean的后面加与不加.item()，在生成onnx文件后，Netron查看网络结构，会产生不一样的网络结构图，详情可以查看我的github仓库里一个网友发的帖子：

https://github.com/hpc203/license-plate-detect-recoginition-opencv/issues/1

在解决这些坑之后，编写了一套使用opencv做车牌检测与识别的程序，包含C++和python两个版本的代码。使用opencv的dnn模块做前向计算，后处理模块是自己使用C++和Python独立编写的。

代码已发布在github上，地址是：https://github.com/hpc203/license-plate-detect-recoginition-opencv

2. opencv与onnxruntime的差异

起初在github上看到一个使用DBNet检测条形码的程序，不过它是基于pytorch框架做的。于是我编写一套程序把pytorch模型转换到onnx文件，使用opencv读取onnx文件做前向计算。编写完程序后在运行时没有出错，但是最后输出的结果跟调用pytorch 的输出结果不一致，并且从可视化结果看，没有检测出图片中的条形码。这时在看到网上有很多使用onnxruntime部署onnx模型的文章，于是决定使用onnxruntime部署，编写完程序后运行，选取几张快递单图片测试，结果如下图所示DBNet检测到的4个点,图中绿色的点，红色的线是把4个连接起来的直线。

并且我还编写了一个函数比较opencv和onnxruntime的输出结果，程序代码和运行结果如下，可以看到在相同输入，读取同一个onnx文件的前提下，opencv和onnxruntime的输出结果竟然不相同。

ONNXRuntime是微软推出的一款推理框架，用户可以非常便利的用其运行一个onnx模型。从这个实验，可以看出相比于opencv库，onnxruntime库对onnx模型支持的更好。

我把这套使用DBNet检测条形码的程序发布在github上，地址是：https://github.com/hpc203/dbnet-barcode

9月2日，我尝试把github上一个很热门的背景抠图程序用opencv部署，但是在opencv读取onnx文件时出错了，但是onnxruntime就能正常读取，代码如下：

import onnxruntime
import numpy as np
import cv2

if __name__=='__main__':
	onnxpath = 'weights/onnx_mobilenetv2_4k.onnx'
	sess = onnxruntime.InferenceSession(onnxpath)

	src = np.random.normal(size=(1, 3, 1080, 1920)).astype(np.float32)
	bgr = np.random.normal(size=(1, 3, 1080, 1920)).astype(np.float32)

	pha, fgr = sess.run(['pha', 'fgr'], {'src': src, 'bgr': bgr})

	net = cv2.dnn.readNet(onnxpath)
	blob = cv2.dnn.blobFromImages([src, bgr])
	net.setInput(blob)
	outs = net.forward(net.getUnconnectedOutLayersNames())
	print(outs)

运行到 net = cv2.dnn.readNet(onnxpath) 这一行时出错了，报错信息如下

 很明显是opencv不能创建TopK这个层，在github仓库里搜索TopK，看到在model/refiner.py里有这么一段代码，截图如下

 它是直接调用pytorch的接口函数topk的，但是这个在opencv的dnn模块里并不支持，可是onnxruntime库读取onnx文件却是正常的，这又说明了onnxruntime库对onnx模型支持的更好。

在10月6日，我在github上发布了一套使用ONNXRuntime部署鲁棒性视频抠图的程序，依然是包含C++和Python两种版本的程序。起初，我想使用opencv的dnn模块作为推理引擎，但是程序运行到cv2.dnn.readNet(modelpath) 这里时报错，因此使用onnxruntime 作为推理引擎，源码地址是：

https://github.com/hpc203/robustvideomatting-onnxruntime

3.  onnxruntime支持3维池化和3维卷积

在第1节讲到opencv不支持3维池化，那么onnxruntime是否支持呢？接着编写了一个程序探索onnxruntime对3维池化的支持情况，代码和运行结果如下，可以看到程序报错了。

查看nn.MaxPool3d的说明文档，截图如下，可以看到它的输入和输出是5维张量，于是修改上面的代码，把输入调整到5维张量。

代码和运行结果如下，可以看到这时候onnxruntime库能正常读取onnx文件，并且它的输出结果跟pytorch的输出结果相等。

继续实验，把三维池化改作三维卷积，代码和运行结果如下，可以看到平均差异在小数点后11位，可以忽略不计。

在第1节讲到过opencv不支持3维池化，那时候的输入张量是4维的，如果把输入张量改成5维的，那么opencv是否就能进行3维池化计算呢？为此，编写代码，验证这个想法。代码和运行结果如下，可以看到在cv2.dnn.blobFromImage这行代码出错了。

查看cv2.dnn.blobFromImage这个函数的说明文档，截图如下，可以看到它的输入image是4维的，这说明它不支持5维的输入。

经过这一系列的程序实验论证，可以看出onnxruntime库对onnx模型支持的更好。如果深度学习模型有3维池化或3维卷积层，那么在转换到onnx文件后，使用onnxruntime部署深度学习是一个不错的选择。并且，我看到一篇文章《yolov3 cpu推理性能比较-onnx opencv darknet》，里面通过一系列实验，最后得出结论是onnxruntime库的推理速度最快。

4.  onnx动态分辨率输入

不过我在做pytorch导出onnx文件时，还发现了一个问题。在torch.export函数里有一个输入参数dynamic_axes，它表示动态的轴，即可变的维度。假如一个神经网络输入是动态分辨率的，那么需要定义dynamic_axes = {'input': {2: 'height', 3: 'width'}, 'output': {2: 'height', 3: 'width'}}，接下来我编写一个程序来验证，代码和运行结果的截图如下

可以看到，在生成onnx文件后，使用onnxruntime库读取，对输入blob的高增加10个像素单位，在run这一步出错了。使用opencv读取onnx文件，代码和运行结果的截图如下，可以看到依然出错了。

通过这个程序实验，让人怀疑torch.export函数的输入参数dynamic_axes是否真的支持动态分辨率输入的。

以上这些程序实验是我在编写算法应用程序时记录下的一些bug和解决方案的，希望能帮助到深度学习算法开发应用人员少走弯路。

此外，DBNet的官方代码里提供了转换到onnx模型文件，于是我依然编写了一套使用opencv部署DBNet文字检测的程序，依然是包含C++和Python两个版本的代码。官方代码的模型是在ICDAR场景文本检测数据集上训练的，考虑到车牌里也含有文字，我把文章开头展示的汽车图片作为输入，程序检测结果如下，可以看到依然能检测到车牌的4个角点，只是不够准确。如果想要获得准确的角点定位，可以在车牌数据集上训练DBNet。

我把使用opencv部署DBNet文字检测的程序发布在github上，程序依然是包含c++和python两种版本的实现，地址是：

https://github.com/hpc203/dbnet-opencv-cpp-python

5月2日，我在github上传了一套程序，使用BiSeNet做人脸面部解析，里面包含了分别基于pytorch, opencv, onnxruntime三种库的程序实现，并且比较了在调用三种库的输入和输出的差异，结果发现调用 pytorch框架的输出和调用opencv和onnxruntime的输出都不同，而opencv和onnxruntime的输出相同(平均差异在小数点后10位)，这又是一个pytorch转onnx时的一个坑。这套程序的地址是 https://github.com/hpc203/Face-Parsing-pytorch-opencv-onnxruntime