pytorch upsample层到onnx，以及到tensorRT的转换

一、pytorch 实现一个上采样层，代码以下

import torch import torch.nn as nn import torch.nn.functional as F import os import cv2 import numpy as np class TestUpsample(nn.Module): def __int__(self): super(TestUpsample, self).__init__() def forward(self, x): x = nn.Upsample(scale_factor=2, mode="nearest")(x) return x

二、使用测试图片，检查模型输出结果，代码以下：

image = cv2.imread("test.jpg")[:,:,::-1] # 1,3,320,320
image = np.transpose(image, [2, 0, 1]) # CHW to NCHW format
image = np.expand_dims(image, axis=0) # Convert the image to row-major order, also known as
image = np.array(image, dtype=np.float32, order='C') image = image.astype(np.float32) image = torch.from_numpy(image) torch_model = TestUpsample() output = torch_model(image) # 1,3,640,640

三、使用 1.5.0 版本onnx和 1.6.0 版本onnx分别将 upsample 层转换到onnx模型

# onnx 1.5.0 版本的转换 # torch_out = torch.onnx._export(torch_model, image, 'upsample-1.5.0.onnx', verbose=True) # 1,3,640,640 # onnx 1.6.0 版本的转换
torch_out = torch.onnx._export(torch_model, image, 'upsample-1.6.0.onnx', verbose=True, opset_version=11) # np.testing.assert_almost_equal(output.data.cpu().numpy(), torch_out.data.cpu().numpy(), decimal=3)

四、使用 Netron-4.1.0 工具查看两个onnx模型的结构：

1.5.0版本的onnx结构

1.6.0版本的onnx结构： 

 注意：1.6.0版本onnx模型中有个 Constant 空节点，在最下面，万恶之源就在这里

五、onnx转tensorrt的时候，就是这个空节点报错。

六、开发环境总结：

转 tensorrt 失败	转 tensorrt 成功
pytorch 1.3.0	pytorch 1.0
onnx 1.6.0	onnx 1.5.0
tensorrt 7.0	tensorrt 7.0
cuda 10.0	cuda 10.0
cudnn 7.6.5	cudnn 7.6.5

 

 

 

 

 

 

 

 

 

七、好不容易将整个目标检测模型转换到了tensorrt框架下，结果发现tensorrt模型推理速度比pytorch原始模型慢3～5ms