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