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
,公式如下:
o u t p u t _ s i z e = c e i l ( ( i n p u t _ s i z e + 2 ∗ p a d d i n g − k e r n e l _ s i z e ) / s t r i d e ) + 1 \mathbf{output\_size} = ceil(( \mathbf{input\_size} + 2 * \mathbf{padding} - \mathbf{kernel\_size}) / \mathbf{stride})+1 output_size=ceil((input_size+2∗padding−kernel_size)/stride)+1
因此通过input_size
、output_size
反推kernel_size
、stride
、padding
,参考官方源码将padding
设为0,那么可推出去kernel_size
、stride
:
s
t
r
i
d
e
=
f
l
o
o
r
(
i
n
p
u
t
_
s
i
z
e
/
o
u
t
p
u
t
_
s
i
z
e
)
\mathbf{stride} = floor(\mathbf{input\_size} / \mathbf{output\_size})
stride=floor(input_size/output_size)
k
e
r
n
e
l
_
s
i
z
e
=
i
n
p
u
t
_
s
i
z
e
−
(
o
u
t
p
u
t
_
s
i
z
e
−
1
)
∗
s
t
r
i
d
e
\mathbf{kernel\_size} = \mathbf{input\_size}- (\mathbf{output\_size}-1) * \mathbf{stride}
kernel_size=input_size−(output_size−1)∗stride
示例
例如,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
- 更新PyTorch到最新版,一般最新版中ONNX的OP支持应该会更多;
- 按下列代码将日志等级调到最高,逐一分析。
import torch
torch.onnx.export(..., verbose=True, ...)
2. 检测ONNX模型
下载Netron可视化自己的ONNX模型,分析是否与PyTorch模型一致,或者与自己想造的模型一致。多看看resize
、shape
、permute
操作,ONNX对这些需要对tensor切片的操作不是很支持。
3. ONNX2TRT
- 更新onnx-tensorrt库,也就是
libnvonnxparser.so
。下面贴一段TRT的安装步骤:- 安装TRT.
- 编译onnx-tensorrt.
- 将
libnvonnxparser.so
移到TRT的lib文件夹中.
- 按下列代码将日志等级调到最高,逐一分析。
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提供的示例<TRT_root>/samples/python/network_api_pytorch_mnist/sample.py
,与之对照的是<TRT_root>/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了。