训练好的模型需要移植到ncnn、mnn、tensorrt、tnn、openvino等前向推理框架上部署
不管是tensorflow、pytorch、mxnet部署时移植到目标框架,那么先将模型转成onnx这个中间商,再转成目标框架或NPU是明智的选择。
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模型
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')
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
转换成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的输出特征比对欧式距离或余弦距离,确保相同输入和预处理方法。