Pytorch模型转onnx、caffe、TorchScript通用代码，onnx模型推理

前言

训练好的模型需要移植到ncnn、mnn、tensorrt、tnn、openvino等前向推理框架上部署

一、模型转onnx

不管是tensorflow、pytorch、mxnet部署时移植到目标框架，那么先将模型转成onnx这个中间商，再转成目标框架或NPU是明智的选择。

二、使用步骤

1.转成onnx示例代码

model = MobileFaceNet()
model.load_state_dict(torch.load('my_model.pth', map_location='cpu'))
model.eval()   先加载模型进入eval()模式
dummy_input = torch.randn(1, 3, 112, 112)  # 你模型的输入   NCHW
torch.onnx.export(model, dummy_input,'my_model.onnx', export_params=True,verbose=False,
                   input_names['input0'],output_names=['output0']) 
onnx_model = onnx.load('./my_model.onnx')  # load onnx model
onnx.checker.check_model(onnx_model)  # check onnx model
print(onnx.helper.printable_graph(onnx_model.graph))  # print a human readable model

转换好模型后onnx里可能很多胶水OP需要用onnx-simplifier进行model slim

pip install onnx-simplifier
python -m onnxsim my_model.onnx my_model-sim.onnx  #生成一个my_model-sim.onnx模型

2.转成TorchScript示例代码

model = MobileFaceNet()
model.load_state_dict(torch.load('my_model.pth', map_location='cpu'))
model.eval()   先加载模型进入eval()模式
dummy_input = torch.randn(1, 3, 112, 112)  # 你模型的输入   NCHW
ts = torch.jit.trace(model, dummy_input)
ts.save('my_model.pt')

3.转成caffe示例代码

git clone https://github.com/hahnyuan/nn_tools, 下载pytorch2caffe代码

from torch.autograd import Variable  #import Variable
import sys
sys.path.append('./nn_tools-master')
import pytorch_to_caffe  # pytorch_to_caffe
model = MobileFaceNet()
model.load_state_dict(torch.load('my_model.pth', map_location='cpu'))
model.eval()   先加载模型进入eval()模式
name = 'my_model'
input = torch.randn(1, 3, 112, 112)  # 你模型的输入   NCHW
pytorch_to_caffe.trans_net(detect_model, input, name)
pytorch_to_caffe.save_prototxt('{}.prototxt'.format(name))  # 生成prototxt
pytorch_to_caffe.save_caffemodel('{}.caffemodel'.format(name))   # 生成caffemodel

4.onnx模型验证推理示例代码

转换成onnx模型后验证onnx模型能否正常推理及精度损失情况

import onnxruntime as ort

def onnx_inferred_demo(ONNX_MODEL):
    image_path = './test.jpg'
    img = cv2.imread(image_path)
    img = cv2.resize(img, (Height, Width))  #模型输入size
    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)


    img = np.transpose(img, (2, 0, 1))  # HWC->CHW

    ort_session = ort.InferenceSession(ONNX_MODEL)
    input_name = ort_session.get_inputs()[0].name  # 'data'
    outputs0 = ort_session.get_outputs()[0].name
    outputs1 = ort_session.get_outputs()[1].name   # 多输出就往下增加
    print(outputs0, outputs1)

    input_blob = np.expand_dims(img, axis=0).astype(np.float32)  # NCHW

    out = ort_session.run([outputs0, outputs1], input_feed={input_name: input_blob})

    print(out[0].shape, out[1].shape)

模型转换后要想验证模型输出有没错误或精度损失，可以进行pytorch和onnx的输出特征比对欧式距离或余弦距离，确保相同输入和预处理方法。