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

 

 

 

posted @ 2020-05-21 17:25  blackx  阅读(4401)  评论(5编辑  收藏  举报