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