PyTorch2ONNX2TensorRT 踩坑日志

PyTorch2ONNX2TensorRT 踩坑日志

从“用PyTorch写的网络，通过ONNX，使用TensorRT序列化，最终完成模型加速”的全流程踩坑日志。

---

2019/12/07 初版
2019/12/17 更新AdaptivePooling, 找BUG思路
2019/12/27 添加AdaptivePooling示例
2020/01/01 添加VGG16示例链接

---

实验环境

ONNX可以不用安装，对ONNX2TRT没有影响，推荐使用anaconda管理包。

• Ubuntu 16.04
• RTX2080TI, Driver Version: 410.79
• CUDA 10.0
• cudnn 7.6.3 (经测低版本如7.5.0无影响)
• pycuda 2019.1.2
• pytorch 1.3.1
• torchvision 0.4.2
• tensorrt 6.0.1.5
• python 3.6.9
  • 经测ONNX无法使用，建议使用python 3.7.x
  • onnx 1.6.0
  • protobuf 3.9.2 (需要降级到3.9.x，不然onnx会报libprotobuf.so.20的错)

---

1. RuntimeError: ONNX export failed: Couldn’t export operator aten::upsample_bilinear2d

无法解决，ONNX2TensorRT报错，待TensorRT后续版本支持，见后文替代方法#4

近似地，应该与警告信息 UserWarning: ONNX export failed on upsample_bilinear2d because align_corners == True not supported 相关联。

原因

转换ONNX使用的版本较低，PyTorch.ONNX不支持。另外，参考源码， torch.onnx.export 默认使用 opset_version=9。

解决办法

警告信息已经完整说明，ONNX's Upsample/Resize operator did not match Pytorch's Interpolation until opset 11.，因此将ONNX的导出代码中规定其版本，具体如下：

import torch
torch.onnx.export(model, ..., opset_version=11)

较完整报错信息

输出的个人信息就被我隐去了，也为了报错、警告的简洁，所以这里叫做“较完整”，此说明后续不再赘述。

UserWarning: You are trying to export the model with onnx:Upsample for ONNX opset version 9. This operator might cause results to not match the expected results by PyTorch.
ONNX's Upsample/Resize operator did not match Pytorch's Interpolation until opset 11. Attributes to determine how to transform the input were added in onnx:Resize in opset 11 to support Pytorch's behavior (like coordinate_transformation_mode and nearest_mode).
We recommend using opset 11 and above for models using this operator. 

UserWarning: ONNX export failed on upsample_bilinear2d because align_corners == True not supported

RuntimeError: ONNX export failed: Couldn't export operator aten::upsample_bilinear2d

---

2. RuntimeError: ONNX export failed: Couldn’t export operator aten::adaptive_avg_pool2d

无法解决，ONNX2TensorRT报错，待TensorRT后续版本支持，见后文替代方法#5

类似错误 aten::adaptive_avg_pool*d：onnx#63, pytorch#14395, discuss.pytorch#30204

原因

因为PyTorch的网络中用了 torch.nn.AdaptiveAvgPool2d ，个人感觉，ONNX没有 avg_pool2d 操作，见ONNX Operator，所以PyTorch.ONNX就会报错 aten::adaptive_avg_pool2d 无法转换。

解决办法

参考pytorch#14395添加额外Option，如下：

import torch
torch.onnx.export(model, ..., operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK)

该方法只是阻止ONNX替换PyTorch的OP、而是使用ATen的OP替换，PyTorch2ONNX能通，但ONNX2TRT却不能通，原因是ONNX phaser识别不到非ONNX的OP。

较完整报错信息

UserWarning: ONNX export failed on adaptive_avg_pool2d because output size that are not factor of input size not supported

RuntimeError: ONNX export failed: Couldn't export operator aten::adaptive_avg_pool2d

---

3. Error: In node 2 (importGather): UNSUPPORTED_NODE: Assertion failed: !(data->getType() == nvinfer1::DataType::kINT32 && nbDims == 1) && “Cannot perform gather on a shape tensor!”

原因

"Cannot perform gather on a shape tensor!"，网络内部使用x_size = x.size()[1:]等类似操作，TensorRT在trace的时候，会被解析成一个shape layer的输出，获得一个shape tensor，用Netron工具可视化就可以发现，对应的node 2实际上是个Constant node，与预期不符。

解决办法

不使用该操作，另一种解法来自onnx-tensorrt#192

x_size = torch.tensor(x.shape)[1:]

---

4. Error: In node 1 (importUpsample): UNSUPPORTED_NODE: Assertion failed: (nbDims >= 1) && (nbDims <= 3)

使用Netron工具可视化模型，找到对应的node 1，就可以发现对应的是F.interpolate(x, size=(128, 128), mode='bilinear', align_corners=False)操作。

原因

目前ONNX2TRT的转换过程中，貌似不支持F.interpolate的bilinear模式，只支持linear和nearest。

解决办法

将所有的bilinear模式替换为nearest模式。

---

5. 使用AvgPooling替换AdaptivePooling

针对2. RuntimeError: ONNX export failed: Couldn't export operator aten::adaptive_avg_pool2d问题，使用AvgPooling替换AdaptivePooling。因为ONNX支持AvgPooling，PyTorch2ONNX、ONNX2TRT流程能够跑通。

原因

目前PyTorch2ONNX流程中，ONNX并不支持AdaptivePooling操作，该操作仅存于PyTorch中。

解决方法

参考[开发技巧]·AdaptivePooling与Max/AvgPooling相互转换一文、PyTorch官方文档可知，AdaptivePooling可通过输入大小input_size自适应控制输出大小output_size，而一般的AvgPooling/MaxPooling则是通过kernel_size、stride、padding来计算output_size，公式如下：

$$\mathbf{o}\mathbf{u}\mathbf{t}\mathbf{p}\mathbf{u}\mathbf{t}\_\mathbf{s}\mathbf{i}\mathbf{z}\mathbf{e}=ceil((\mathbf{i}\mathbf{n}\mathbf{p}\mathbf{u}\mathbf{t}\_\mathbf{s}\mathbf{i}\mathbf{z}\mathbf{e}+2\ast \mathbf{p}\mathbf{a}\mathbf{d}\mathbf{d}\mathbf{i}\mathbf{n}\mathbf{g}-\mathbf{k}\mathbf{e}\mathbf{r}\mathbf{n}\mathbf{e}\mathbf{l}\_\mathbf{s}\mathbf{i}\mathbf{z}\mathbf{e})/\mathbf{s}\mathbf{t}\mathbf{r}\mathbf{i}\mathbf{d}\mathbf{e})+1$$

因此通过input_size、output_size反推kernel_size、stride、padding，参考官方源码将padding设为0，那么可推出去kernel_size、stride：

$$\mathbf{s}\mathbf{t}\mathbf{r}\mathbf{i}\mathbf{d}\mathbf{e}=floor(\mathbf{i}\mathbf{n}\mathbf{p}\mathbf{u}\mathbf{t}\_\mathbf{s}\mathbf{i}\mathbf{z}\mathbf{e}/\mathbf{o}\mathbf{u}\mathbf{t}\mathbf{p}\mathbf{u}\mathbf{t}\_\mathbf{s}\mathbf{i}\mathbf{z}\mathbf{e})$$
$$\mathbf{k}\mathbf{e}\mathbf{r}\mathbf{n}\mathbf{e}\mathbf{l}\_\mathbf{s}\mathbf{i}\mathbf{z}\mathbf{e}=\mathbf{i}\mathbf{n}\mathbf{p}\mathbf{u}\mathbf{t}\_\mathbf{s}\mathbf{i}\mathbf{z}\mathbf{e}-(\mathbf{o}\mathbf{u}\mathbf{t}\mathbf{p}\mathbf{u}\mathbf{t}\_\mathbf{s}\mathbf{i}\mathbf{z}\mathbf{e}-1)\ast \mathbf{s}\mathbf{t}\mathbf{r}\mathbf{i}\mathbf{d}\mathbf{e}$$

示例

例如，PyTorch网络的某一层含有nn.AdaptiveAvgPool2d(output_size=(14,14))，它的output_size为(14, 14)，该层的输入特征图大小为10*128*128，那么输出的特征图大小为10*14*14，那么带入公式，可计算出nn.AvgPool2d(kernel_size, stride)的stride=(int(128/14), int(128/14)), kernel_size=((128-(14-1)*stride, (128-(14-1)*stride)，验证如下：

import torch
from   torch import nn

input = torch.randn(10, 36, 36)
AAVP = nn.AdaptiveAvgPool2d(output_size=(12,12))
AVP  = nn.AvgPool2d(kernel_size=(3,3), stride=(3,3))

output_AAVP = AAVP(input)
output_AVP  = AVP(input)

---

6. PyTorch2ONNX、ONNX2TRT到底哪里出了问题？

下面是遇到无法解决的问题后该找谁问的一个思路：

PyTorch2ONNX是调用的PyTorch内部的转换脚本，所以PyTorch2ONNX出了问题，需要去PyTorch那里的issue寻找解决办法；ONNX2TRT是使用ONNX自己写的转换脚本onnx-tensorrt，同理如果ONNX2TRT出了问题，就需要到它的那里找解决办法；在产生好TRT模型后，使用TRT模型进行推理出问题了，那就要去TRT那里问了，有GitHub和官方论坛可以使用。

那怎么让报错暴露出来呢，下面是一些办法。

解决方法

按下列方法多半能找到问题所在。

1. PyTorch2ONNX

1. 更新PyTorch到最新版，一般最新版中ONNX的OP支持应该会更多；
2. 按下列代码将日志等级调到最高，逐一分析。

import torch
torch.onnx.export(..., verbose=True, ...)

2. 检测ONNX模型

下载Netron可视化自己的ONNX模型，分析是否与PyTorch模型一致，或者与自己想造的模型一致。多看看resize、shape、permute操作，ONNX对这些需要对tensor切片的操作不是很支持。

3. ONNX2TRT

1. 更新onnx-tensorrt库，也就是libnvonnxparser.so。下面贴一段TRT的安装步骤：
   1. 安装TRT.
   2. 编译onnx-tensorrt.
   3. 将libnvonnxparser.so移到TRT的lib文件夹中.
2. 按下列代码将日志等级调到最高，逐一分析。

import tensorrt as trt
TRT_LOGGER = trt.Logger(trt.Logger.VERBOSE)

---

最终解决办法

放弃ONNX2TRT吧，PyTorch与ONNX与TRT的版本难以互相支持，在版本的迭代中任意节点不支持了，整个链路就会断掉，另外TRT是闭源的项目，你完全不知道ONNX2TRT的过程中出了哪些问题，就算有堆栈信息，也不可能根据信息去trace它的错误。所以，直接使用TRT提供的api直接构建网络，是最复杂、也是最简单直接的方法。

Pytorch 2 TRT python API

使用TRT提供的python接口，构建网络，整个流程十分简单，大家可以看看TRT提供的示例/samples/python/network_api_pytorch_mnist/sample.py，与之对照的是/samples/python/network_api_pytorch_mnist/model.py：

def populate_network(network, weights):
    # Configure the network layers based on the weights provided.
    input_tensor = network.add_input(name=ModelData.INPUT_NAME, dtype=ModelData.DTYPE, shape=ModelData.INPUT_SHAPE)
    """
    TRT python API
    """
    network.mark_output(tensor=fc2.get_output(0))

你只需要把这个populate_network写出来就好了，weights就是网络的权重了，由torch.load()得到，是不是超级简单啊。想使用PyTorch的F.interpolate的bilinear模式？TRT提供！下篇日志将会记录“如何使用TRT python API搭建简单的VGG16网络”，我再也不想用ONNX2TRT了。