Torch模型转换

背景

接上一篇部署，最近由于需要移植之前在Jestson设备上部署过的模型，所以重新做了一次模型转换和模型部署的流程，正好一边温习一边记录。

Torch to ONNX

ONNX，全称Open Neural Network Exchange，奇奇怪怪的名字不重要，缩写好记就行。简单理解就是一个中间表达，把一些常用的如torch这样的框架统一用中间框架来表达，后续的轻量级框架设计只需要对接ONNX做算子转换即可。
torch完成训练验证后，保存的模型直接加载部署过于笨重，在生产环境占用资源多，推理速度慢是常见的事情。因此一些高性能推理框架就应运而生，但辛辛苦苦写好训练代码，总不能出一个新框架就复写一次，于是乎ONNX就用上了。
官方的ONNX不断迭代，把一些新的或复杂的操作更新到算子转换库里，这样就不需要我们手动实现从torch到ONNX的算子转换代码。转换过程可以理解为把模型的运行过程提取出来，导出一个计算图，对图中每一个节点操作做替换。导出计算图的方案有两种，基于tracing的和基于script的，区别的话就是tracing实际地跑了一次模型，记录这次走过的路径图，只对该图的操作节点做转换；而script则是解析整个模型的计算图，依次转换。由此，显而易见的，tracing方式不支持forward()中存在的控制流程，但可惜的是script方式总存在一些奇怪的问题，对于tracing方式的支持似乎更好。
得到计算图后，对图中的节点做操作映射，映射为ONNX的表达方式，最终即可得到转换后的ONNX模型。

def transe2onnx(self, model):
        batch_size = 1 								# 批处理大小
        input_shape = (3, 360, 640) 				# 输入数据  

        x = torch.randn(batch_size,*input_shape).cuda() 	# 生成张量
        export_onnx_file = "test.onnx" 				# 目的ONNX文件名
        torch.onnx.export(model,
                    x,
                    export_onnx_file,
                    opset_version=11,
                    verbose=False)

对于静态图，构造一个维度与模型输入相同的tensor即可，需要注意的就是操作集的版本，有的操作只在某一版本实现转换。

ONNX to TensorRT

TensorRT，Nvidia的高性能推理框架，用就完事儿了。
目前了解到有两种方式将torch的模型转为tensorRT模型，其一是torch2trt这么一个工具，可以对一些基本操作构造的torch模型做转换，简单直接好用，但仅限于常用算子。其二则是将torch先转ONNX，再把ONNX转到tensorRT，略繁琐，但适用性强。
ONNX转tensorRT又有两种转换方式，其一是安装好tensorRT后自带了一个脚本工具trtexec，其二是自己写脚本，本质上感觉没有太大区别，自己写脚本可以指定一些细节，而trtexec则固定了几类参数供选择。
使用trtexec需要注意的是动态输入数据的情况，固定维度的输入无需考虑太多，而可变维度需要配置-–minShapes，-–maxShapes来规定可变维度的范围，–-optShapes则规定了本次转换过程使用的维度。
基本上只要做好了第一步的转换，得到正确的ONNX，第二步转tensorRT问题都不会太多。

问题

具体实施起来遇到的问题特别多，比如维度啦、算子版本啦之类的，总之是什么问题都会有。
解决过程中记录了几个比较有意思的点：
第一个是给pytorch加自定义的C++扩展，原理和解法在官方博文中找到了一些，包括把自定义算法搞成一个库打包安装什么的，头一次接触这个，还是挺有意思的。
第二个是精度问题，做模型量化单精度转半精度的时候遇到了贼奇怪的问题，猜测是精度溢出，目前也没有彻底解决，只对一些能改的操作做了修改，是一个需要注意有坑的点。
其余问题大多能在网上找到对应的解法，转ONNX这里如果模型用到了比较多的新的或者复杂的aten操作，是需要做一些替换或者改写的，如果实在改不了就需要手动自定义加映射了，目前还没有具体做到自定义映射实现这步，感觉是有点难度和价值的东西。
部署的事情属于比较杂的，需要解决的问题也五花八门，但只要耐心找，解法总能找到，从原理中、从博客中或者从盖的楼中，愿诸君的部署过程武运昌隆。