OpenMMlab导出swin-transformer模型并用onnxruntime和tensorrt推理
导出onnx文件
通过mmpretrain 导出swin-transformer的onnx文件非常容易，注意需设置 opset_version=12这里是一个坑，刚开始设置的opset_version=11后续转换trtengine的时候会有问题。
import torch
from mmpretrain import get_model, inference_model
model = get_model('swin-tiny_16xb64_in1k',pretrained='swin_tiny_224_b16x64_300e_imagenet_20210616_090925-66df6be6.pth', device='cpu')
input = torch.zeros(1, 3, 224, 224)
out = model(input)
print(torch.argmax(out, dim=1))
torch.onnx.export(model, input, "swin_transformer.onnx", opset_version=12)
如果安装了mmdeploy，也可以通过下面的脚本进行模型转换：
from mmdeploy.apis import torch2onnx
from mmdeploy.backend.sdk.export_info import export2SDK
img = 'demo.JPEG'
work_dir = './work_dir/onnx/swin_transformer'
save_file = './end2end.onnx'
deploy_cfg = 'mmdeploy/configs/mmpretrain/classification_onnxruntime_dynamic.py'
model_cfg = 'mmpretrain/configs/swin_transformer/swin-tiny_16xb64_in1k.py'
model_checkpoint = './checkpoints/swin_tiny_224_b16x64_300e_imagenet_20210616_090925-66df6be6.pth'
device = 'cpu'
# 1. convert model to onnx
torch2onnx(img, work_dir, save_file, deploy_cfg, model_cfg, model_checkpoint, device)
# 2. extract pipeline info for sdk use (dump-info)
export2SDK(deploy_cfg, model_cfg, work_dir, pth=model_checkpoint, device=device)
onnxruntime推理
python推理：
import cv2
import numpy as np
import onnxruntime
use_letterbox = True
input_shape = (224, 224)
def letterbox(im, new_shape=(416, 416), color=(114, 114, 114)):
# Resize and pad image while meeting stride-multiple constraints
shape = im.shape[:2] # current shape [height, width]
# Scale ratio (new / old)
r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
# Compute padding
new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
dw, dh = (new_shape[1] - new_unpad[0])/2, (new_shape[0] - new_unpad[1])/2 # wh padding
top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))
left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
if shape[::-1] != new_unpad: # resize
im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) # add border
return im
if __name__ == '__main__':
img = cv2.imread('goldfish.jpg')
if use_letterbox:
or_img = letterbox(img, input_shape)
else:
or_img = cv2.resize(img, input_shape)
img = or_img[:,:,::-1].transpose(2,0,1) #BGR2RGB和HWC2CHW
img = img.astype(dtype=np.float32)
img[0,:] = (img[0,:] - 123.675) / 58.395
img[1,:] = (img[1,:] - 116.28) / 57.12
img[2,:] = (img[2,:] - 103.53) / 57.375
img = np.expand_dims(img, axis=0)
onnx_session = onnxruntime.InferenceSession("swin_transformer.onnx", providers=['CUDAExecutionProvider','CPUExecutionProvider'])
input_name=[]
for node in onnx_session.get_inputs():
input_name.append(node.name)
output_name=[]
for node in onnx_session.get_outputs():
output_name.append(node.name)
input_feed={}
for name in input_name:
input_feed[name] = img
pred = onnx_session.run(None, input_feed)[0]
print(pred)
print(np.argmax(pred))
C++推理：
#include 如果安装了mmdeploy，也可以这样进行推理：
python推理：
import cv2
from mmdeploy_runtime import Classifier
img = cv2.imread('mmpretrain/demo/demo.JPEG')
classifier = Classifier(model_path='work_dir/onnx/swin_transformer', device_name='cpu')
result = classifier(img)
for label_id, score in result:
print(label_id, score)
C++推理可以参考：https://github.com/open-mmlab/mmdeploy/blob/main/demo/csrc/c/image_classification.cpp
导出engine文件
这里通过trtexec转换onnx文件，LZ的版本是TensorRT-8.2.1.8。
需要先使用onnxsim简化模型，这里是第二个坑，否则会报错。
import onnx
from onnxsim import simplify
onnx_model = onnx.load("swin_transformer.onnx") # load onnx model
model_simp, check = simplify(onnx_model)
assert check, "Simplified ONNX model could not be validated"
onnx.save(model_simp, "swin_transformer_sim.onnx")
再去ensorRT的bin目录下运行
./trtexec.exe --onnx=swin_transformer_sim.onnx --saveEngine=swin_transformer.engine --workspace=20480
第三个坑是如果不加上–workspace参数可能会因内存不足报错，LZ的机器有32G内存索性就设了20G的工作空间，可以根据自己的内存大小酌情设置该参数。
至此，不出意外可以成功导出engine文件。
mmdeploy导出LZ没有尝试成功，不知道是环境配置问题还是在windows系统下库的bug。
tensorrt推理
python推理：（待补充）
C++推理：
// tensorRT include
#include