实操教程｜PyTorch转ONNX详解

之前几个月参与了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都变成一个层。。。）

tensor i = funcA();
if(i==0)
  j = funcB(i);
else
  j = funcC(i);
funcD(j);

tensor i = funcA();
coarse_func(tensor i) {
  if(i==0) return funcB(i);
  else return funcC(i);
}
funcD(coarse_func(i))