Pytorch转化ONNX过程代码实操
本文作者总结了自己参与pytorch到onnx的模型转换转换工作中的经验，主要介绍了该转换工作的意义，模型部署的路径以及pytorch本身的局限。
之前几个月参与了openmmlab的模型转onnx的工作（github account: drcut），主要目标是支持openmmlab的一些模型从pytorch到onnx的转换。这几个月虽然没做出什么成果，但是踩了很多坑，在这里记录下来，希望可以帮助其他人。
这篇是第一部分，理论篇，主要介绍了和代码无关的一些宏观问题。再接下来我会专门写一篇实战篇，针对openmmlab中一些具体代码做分析，说明pytorch转化onnx过程中的一些代码上的技巧和注意事项。
（1）pytorch转onnx的意义
一般来说转onnx只是一个手段，在之后得到onnx模型后还需要再将它做转换，比如转换到tensorrt上完成部署，或者有的人多加一步，从onnx先转换到caffe，再从caffe到tensorrt。原因是caffe对tensorrt更为友好，这里关于友好的定义后面会谈。
因此在转onnx工作开展之前，首先必须明确目标后端。onnx只是一个格式，就和json一样。只要你满足一定的规则，都算是合法的，因此单纯从pytorch转成一个onnx文件很简单。但是不同后端设备接受的onnx是不一样的，因此这才是坑的来源。
pytorch自带的torch.onnx.export转换得到的onnx，onnxruntime需要的onnx，tensorrt需要的onnx都是不同的。
这里面举一个最简单的maxpool的例：
maxunpool可以被看作maxpool的逆运算，咱们先来看一个maxpool的例子，假设有如下一个c*h*w的tensor（shape[2, 3, 3])，其中每个channel的二维矩阵都是一样的，如下所示
在这种情况下，如果我们在pytorch对它调用maxpool（kernel_size=2, stride=1,pad=0）
那么会得到两个输出，第一个输出是maxpool之后的值：
另一个是maxpool的idx，即每个输出对应原来的哪个输入，这样做反向传播的时候就可以直接把输出的梯度传给对应的输入：
细心的同学会发现其实maxpool的idx还可以有另一种写法：
，
即每个channel的idx放到一起，并不是每个channel单独从0开始。这两种写法都没什么问题，毕竟只要反向传播的时候一致就可以。
但是当我在支持openmmediting的时候，会涉及到maxunpool，即maxpool的逆运算：输入maxpoolid和maxpool的输出，得到maxpool的输入。
pytorch的maxunpool实现是接收每个channel都从0开始的idx格式，而onnxruntime则相反。因此如果你希望用onnxruntime跑一样的结果，那么必须对输入的idx（即和pytorch一样的输入）做额外的处理才可以。换言之，pytorch转出来的神经网络图和onnxruntime需要的神经网络图是不一样的。
（2）onnx与caffe
主流的模型部署有两种路径，以tensorrt为例，一种是pytorch->onnx->tensorrt，另一种是pytorch->caffe->tensorrt。个人认为目前后者更为成熟，这主要是onnx，caffe和tensorrt的性质共同决定的
上面的表列了onnx和caffe的几点区别，其中最重要的区别就是op的粒度。举个例子，如果对bert的attention层做转换，onnx会把它变成matmul，scale，softmax的组合，而caffe可能会直接生成一个叫做multi-head attention的层，同时告诉cuda工程师：“你去给我写一个大kernel“（很怀疑发展到最后会不会把resnet50都变成一个层。。。）
因此如果某天一个研究员提了一个新的state-of-the-art的op，很可能它直接就可以被转换成onnx（如果这个op在pytorch的实现全都是用aten的库拼接的），但是对于caffe的工程师，需要重新写一个kernel。 细粒度op的好处就是非常灵活，坏处就是速度会比较慢。这几年有很多工作都是在做op fushion（比如把卷积和它后面的relu合到一起算），xla和tvm都有很多工作投入到了op fushion，也就是把小op拼成大op。 tensorrt是nvidia推出的部署框架，自然性能是首要考量的，因此他们的layer粒度都很粗。在这种情况下把caffe转换过去有天然的优势。 除此之外粗粒度也可以解决分支的问题。tensorrt眼里的神经网络就是一个单纯的dag：给定固定shape的输入，执行相同的运算，得到固定shape的输出。 **目前tensorrt的一个发展方向是支持dynamic shape，但是还很不成熟。
tensor i = funca(); if(i==0) j = funcb(i); else j = funcc(i); funcd(j); 对于上面的网络，假设funca,funcb,funcc和funcd都是onnx支持的细粒度算子，那么onnx就会面临一个困难，它转换得到的dag要么长这样：funca->funcb->funcd，要么funca->funcc->funcd。但是无论哪种肯定都是有问题的。 而caffe可以用粗粒度绕开这个问题tensor i = funca(); coarse_func(tensor i) { if(i==0) return funcb(i); else return funcc(i); } funcd(coarse_func(i)) 因此它得到的dag是：funca->coarse_func->funcd 当然，caffe的代价就是苦逼的hpc工程师就要手写一个coarse_func kernel。。。（希望deep learning compiler可以早日解放hpc工程师）
(3)pytorch本身的局限 熟悉深度学习框架的同学都知道，pytorch之所以可以在tensorflow已经占据主流的情况下横空出世，成功抢占半壁江山，主要的原因是它很灵活。举个不恰当的例子，tensorflow就像是c++，而pytorch就是python。 tensorflow会把整个神经网络在运行前做一次编译，生成一个dag（有向无环图），然后再去跑这张图。pytorch则相反，属于走一步看一步，直到运行到这个节点算出结果，才知道下一个节点该算啥。 onnx其实就是把上层深度学习框架中的网络模型转换成一张图，因为tensorflow本身就有一张图，因此只需要直接把这张图拿到手，修修补补就可以。 但是对于pytorch，没有任何图的概念，因此如果想完成pytorch到onnx的转换，就需要让onnx再旁边拿个小本子，然后跑一遍pytorch，跑到什么就把什么记下来，把记录的结果抽象成一张图。因此pytorch转onnx有两个天然的局限。 1. 转换的结果只对特定的输入。如果换一个输入导致网络结构发生了变化，onnx是无法察觉的（最常见的情况是如果网络中有if语句，这次的输入走了if的话，onnx就只会生成if对应的图，把else里面全部的信息都丢掉）。 2. 需要比较多的计算量，因为需要真刀真枪的跑一遍神经网络。 ps：针对于以上的两个局限，我的本科毕设论文提出了一种解决方案，就是通过编译器里面的词法分析，语法分析直接扫描pytorch或者tensorflow的源代码得到图结构，这样可以轻量级的完成模型到onnx的转换，同时也可以得到分支判断等信息，这里放一个github链接(https://github.com/drcut/nn_transform)，希望大家多多支持 *目前pytorch官方希望通过用torchscript的方式解决分支语句的问题，但据我所知还不是很成熟。
迅为IMX6UL核心板在便携式医疗设备中的应用方案
中国长城总裁一行莅临易华录调研
中国互联网TOP100榜单公布：VIPKID成最具价值教育企业
一文解读燃料电池的优势及发展
分析云网络、能源基础设施和工业自动化市场的变化发展趋势
Pytorch转化ONNX过程代码实操
台积电：建议汽车芯片转向先进节点
索尼蓝牙音箱SRS-XB20评测 值不值得买
未来已来:网络空间安全迎接2.0时代
NVIDIA Ampere架构解析：相比上一代做了哪些提升
你不知道的五种距离传感器！
如何构建弹性、高可用的微服务？
电影中的物联网是如何想象的
阿里巴巴如何走好自己的自动驾驶之路？
第3届亚太国际智能装备博览会圆满落幕，台湾高技展会后硕果累累
电子管管脚如何识别
奇葩说版米粉节，小米要你先看用户辩论！
STM32如何配置外部中断？STM32外部中断配置步骤
艾比森LED创意从机场 到比赛现场LED大屏无处不美！
MEEP：打造未来开放式HPC生态系统的数字实验室
之前几个月参与了openmmlab的模型转onnx的工作（github account: drcut），主要目标是支持openmmlab的一些模型从pytorch到onnx的转换。这几个月虽然没做出什么成果，但是踩了很多坑，在这里记录下来，希望可以帮助其他人。
这篇是第一部分，理论篇，主要介绍了和代码无关的一些宏观问题。再接下来我会专门写一篇实战篇，针对openmmlab中一些具体代码做分析，说明pytorch转化onnx过程中的一些代码上的技巧和注意事项。
（1）pytorch转onnx的意义
一般来说转onnx只是一个手段，在之后得到onnx模型后还需要再将它做转换，比如转换到tensorrt上完成部署，或者有的人多加一步，从onnx先转换到caffe，再从caffe到tensorrt。原因是caffe对tensorrt更为友好，这里关于友好的定义后面会谈。
因此在转onnx工作开展之前，首先必须明确目标后端。onnx只是一个格式，就和json一样。只要你满足一定的规则，都算是合法的，因此单纯从pytorch转成一个onnx文件很简单。但是不同后端设备接受的onnx是不一样的，因此这才是坑的来源。
pytorch自带的torch.onnx.export转换得到的onnx，onnxruntime需要的onnx，tensorrt需要的onnx都是不同的。
这里面举一个最简单的maxpool的例：
maxunpool可以被看作maxpool的逆运算，咱们先来看一个maxpool的例子，假设有如下一个c*h*w的tensor（shape[2, 3, 3])，其中每个channel的二维矩阵都是一样的，如下所示
在这种情况下，如果我们在pytorch对它调用maxpool（kernel_size=2, stride=1,pad=0）
那么会得到两个输出，第一个输出是maxpool之后的值：
另一个是maxpool的idx，即每个输出对应原来的哪个输入，这样做反向传播的时候就可以直接把输出的梯度传给对应的输入：
细心的同学会发现其实maxpool的idx还可以有另一种写法：
，
即每个channel的idx放到一起，并不是每个channel单独从0开始。这两种写法都没什么问题，毕竟只要反向传播的时候一致就可以。
但是当我在支持openmmediting的时候，会涉及到maxunpool，即maxpool的逆运算：输入maxpoolid和maxpool的输出，得到maxpool的输入。
pytorch的maxunpool实现是接收每个channel都从0开始的idx格式，而onnxruntime则相反。因此如果你希望用onnxruntime跑一样的结果，那么必须对输入的idx（即和pytorch一样的输入）做额外的处理才可以。换言之，pytorch转出来的神经网络图和onnxruntime需要的神经网络图是不一样的。
（2）onnx与caffe
主流的模型部署有两种路径，以tensorrt为例，一种是pytorch->onnx->tensorrt，另一种是pytorch->caffe->tensorrt。个人认为目前后者更为成熟，这主要是onnx，caffe和tensorrt的性质共同决定的
上面的表列了onnx和caffe的几点区别，其中最重要的区别就是op的粒度。举个例子，如果对bert的attention层做转换，onnx会把它变成matmul，scale，softmax的组合，而caffe可能会直接生成一个叫做multi-head attention的层，同时告诉cuda工程师：“你去给我写一个大kernel“（很怀疑发展到最后会不会把resnet50都变成一个层。。。）
因此如果某天一个研究员提了一个新的state-of-the-art的op，很可能它直接就可以被转换成onnx（如果这个op在pytorch的实现全都是用aten的库拼接的），但是对于caffe的工程师，需要重新写一个kernel。 细粒度op的好处就是非常灵活，坏处就是速度会比较慢。这几年有很多工作都是在做op fushion（比如把卷积和它后面的relu合到一起算），xla和tvm都有很多工作投入到了op fushion，也就是把小op拼成大op。 tensorrt是nvidia推出的部署框架，自然性能是首要考量的，因此他们的layer粒度都很粗。在这种情况下把caffe转换过去有天然的优势。 除此之外粗粒度也可以解决分支的问题。tensorrt眼里的神经网络就是一个单纯的dag：给定固定shape的输入，执行相同的运算，得到固定shape的输出。 **目前tensorrt的一个发展方向是支持dynamic shape，但是还很不成熟。
tensor i = funca(); if(i==0) j = funcb(i); else j = funcc(i); funcd(j); 对于上面的网络，假设funca,funcb,funcc和funcd都是onnx支持的细粒度算子，那么onnx就会面临一个困难，它转换得到的dag要么长这样：funca->funcb->funcd，要么funca->funcc->funcd。但是无论哪种肯定都是有问题的。 而caffe可以用粗粒度绕开这个问题tensor i = funca(); coarse_func(tensor i) { if(i==0) return funcb(i); else return funcc(i); } funcd(coarse_func(i)) 因此它得到的dag是：funca->coarse_func->funcd 当然，caffe的代价就是苦逼的hpc工程师就要手写一个coarse_func kernel。。。（希望deep learning compiler可以早日解放hpc工程师）
(3)pytorch本身的局限 熟悉深度学习框架的同学都知道，pytorch之所以可以在tensorflow已经占据主流的情况下横空出世，成功抢占半壁江山，主要的原因是它很灵活。举个不恰当的例子，tensorflow就像是c++，而pytorch就是python。 tensorflow会把整个神经网络在运行前做一次编译，生成一个dag（有向无环图），然后再去跑这张图。pytorch则相反，属于走一步看一步，直到运行到这个节点算出结果，才知道下一个节点该算啥。 onnx其实就是把上层深度学习框架中的网络模型转换成一张图，因为tensorflow本身就有一张图，因此只需要直接把这张图拿到手，修修补补就可以。 但是对于pytorch，没有任何图的概念，因此如果想完成pytorch到onnx的转换，就需要让onnx再旁边拿个小本子，然后跑一遍pytorch，跑到什么就把什么记下来，把记录的结果抽象成一张图。因此pytorch转onnx有两个天然的局限。 1. 转换的结果只对特定的输入。如果换一个输入导致网络结构发生了变化，onnx是无法察觉的（最常见的情况是如果网络中有if语句，这次的输入走了if的话，onnx就只会生成if对应的图，把else里面全部的信息都丢掉）。 2. 需要比较多的计算量，因为需要真刀真枪的跑一遍神经网络。 ps：针对于以上的两个局限，我的本科毕设论文提出了一种解决方案，就是通过编译器里面的词法分析，语法分析直接扫描pytorch或者tensorflow的源代码得到图结构，这样可以轻量级的完成模型到onnx的转换，同时也可以得到分支判断等信息，这里放一个github链接(https://github.com/drcut/nn_transform)，希望大家多多支持 *目前pytorch官方希望通过用torchscript的方式解决分支语句的问题，但据我所知还不是很成熟。
迅为IMX6UL核心板在便携式医疗设备中的应用方案
中国长城总裁一行莅临易华录调研
中国互联网TOP100榜单公布：VIPKID成最具价值教育企业
一文解读燃料电池的优势及发展
分析云网络、能源基础设施和工业自动化市场的变化发展趋势
Pytorch转化ONNX过程代码实操
台积电：建议汽车芯片转向先进节点
索尼蓝牙音箱SRS-XB20评测 值不值得买
未来已来:网络空间安全迎接2.0时代
NVIDIA Ampere架构解析：相比上一代做了哪些提升
你不知道的五种距离传感器！
如何构建弹性、高可用的微服务？
电影中的物联网是如何想象的
阿里巴巴如何走好自己的自动驾驶之路？
第3届亚太国际智能装备博览会圆满落幕，台湾高技展会后硕果累累
电子管管脚如何识别
奇葩说版米粉节，小米要你先看用户辩论！
STM32如何配置外部中断？STM32外部中断配置步骤
艾比森LED创意从机场 到比赛现场LED大屏无处不美！
MEEP：打造未来开放式HPC生态系统的数字实验室