OpenMMlab导出mobilenet-v2的onnx模型并推理-CSDN博客

使用mmpretrain导出mobilenet-v2的onnx模型:

import torch
import numpy as np
from mmpretrain import get_model

model = get_model('mobilenet-v2_8xb32_in1k',pretrained='mobilenet_v2_batch256_imagenet_20200708-3b2dc3af.pth', device='cpu') 
input = torch.zeros(1, 3, 224, 224)
out = model(input)
print(torch.argmax(out, dim=1))
torch.onnx.export(model, input, "mobilenet-v2.onnx", opset_version=11)

或者安装有mmdeploy的话可以通过如下方法导出

from mmdeploy.apis import torch2onnx
from mmdeploy.backend.sdk.export_info import export2SDK

img = 'demo.JPEG'
work_dir = './work_dir/onnx/mobilenet_v2'
save_file = './end2end.onnx'
deploy_cfg = 'mmdeploy/configs/mmpretrain/classification_onnxruntime_static.py'
model_cfg = 'mmpretrain/configs/mobilenet_v2/mobilenet-v2_8xb32_in1k.py'
model_checkpoint = './checkpoints/mobilenet_v2_batch256_imagenet_20200708-3b2dc3af.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进行推理

import cv2
import numpy as np
import onnxruntime


img = cv2.imread('goldfish.jpg')
img = cv2.resize(img, (224,224))
img = 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("mobilenet-v2.onnx", providers=['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)
print(np.argmax(pred))