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

1、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

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

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

3、使用 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)

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

1.5.0版本的onnx结构

1.6.0版本的onnx结构： 

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

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

6、开发环境总结：

转 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

 

 

 

 

 

 

 

 

 

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