Pytorch模型转onnx、caffe、TorchScript通用代码，Ctrl + V直接嵌入

前言

训练好的模型需要移植到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')

2.转成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

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