onnx 模型转换与 onnxruntime 和caffe2 推理速度比较
背景
pytorch 模型通常包括网络结构.py文件和模型参数文件.pth,通常需要转换为onnx格式再向其他终端或移动端部署(onnx模型一般用于中间部署阶段,相当于翻译的作用)。移动端通常使用ncnn、mnn、Caffe2进行inference,实现实时计算。
本文展示pytorch转onnx的模型在PC端 以CPU方式运行。
1、模型转换
1.1. pytorch 模型转onnx 模型
def torch2onnx(model, save_path):
"""
:param model:
:param save_path: XXX/XXX.onnx
:return:
"""
model.eval()
data = torch.rand(1, 3, 48, 64)
input_names = ["input"]
output_names = ["cls","lmk"]
torch.onnx._export(model, data, save_path, export_params=True, opset_version=11, input_names=input_names, output_names=output_names)
print("torch2onnx finish.")
在实际使用过程中我们希望inference时支持batch=2的操作,比如左右眼分类。
可以修改
data = torch.rand(2, 3, 48, 64)
这样在onnx 模型inference时必须要输入两张图,运行时间理论上应该能比单张图输入稍快点。以下模型推理均按batch=2实验.
1.2. onnx 模型简化
onnx 模型简化是可以理解为,根据输入图像尺寸对pytorch动态图中动态变化的量进行固定,比如Gather取tensor维度、unsqueeze/squeeze、shape等算子,其参数与数据流尺寸相关,当输入图像尺寸固定后,相关参数可以转换为constant 值。左图是原始onnx模型图,右侧是简化后的图。可以看到图中数据流尺寸已经很具体,且相关的shape,取维度操作均被干掉了。
安装:
pip install onnx-simplifier
默认安装最新版,也可以指定版本,因为某些情况下依赖库版本不匹配
pip install onnx-simplifier==0.2.18
使用:
在pytorch 转换为onnx模型之后,执行:
python -m onnxsim ${model} ${model}
如果需要在python中调用,执行:
import subprocess
subprocess.call('python -m onnxsim {} {}'.formal(out_model_name,out_model_name) , shell=True)
2、PC端运行onnx模型
2.1.预处理
加载图像和预处理:
import cv2
import numpy as np
img=cv2.imread("test.png")
img=cv2.resize(img,(224,224),interpolation=cv2.INTER_AREA)
img2=img.copy().astype(np.float32)
img2=(img2/255-0.5)/0.5 # 归一化
img_array=np.expand_dims(img2.transpose((2,0,1)),axis=0) #BGR->RGB
img_array=img_array.repeat(2,axis=0) #2 batch data in
当然,对于pytorch 也可以写成:
import torchvision.transforms as transforms
img = cv2.imread('test.png')
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
tensor = transforms.ToTensor()(img)
tensor = transforms.Normalize((0.5,0.5,0.5),(0.5,0.5,0.5))(tensor)
tensor = tensor.unsqueeze_(0)
img_array=tensor.cpu().numpy()
2.2.onnxruntime 运行导出的模型
可参考:帮助文档
安装:
pip install onnxruntime
pip install onnxruntime-gpu
使用:
import onnxruntime
session=onnxruntime.InferenceSession('out.onnx')
input={"input.1":img_array.astype(np.float32)}
output=session.run(None,input)
运行10次:
平均运行时间>550ms
2.3.caffe2 运行导出的模型
import onnx
import caffe2.python.onnx.backend
modelFile=onnx.load('out.onnx')
有两种写法:
- run_model 加载并运行模型
output=caffe2.python.onnx.backend.run_model( modelFile, img_array.astype(np.float32))
运行10次:
平均时间>130ms
- 先加载模型后运行
prepared_backend=caffe2.python.onnx.backend.prepare(modelFile)
W={modelFile.graph.input[0].name:img_array.astype(np.float32)}
output=prepared_backend.run(W)
运行10次:
平均时间约11ms 左右,可见加载模型占了较长的时间。
2.4.总结:
onnx 模型在caffe2 CPU上运行时间远小于onnxruntime的方式。
此外还有Tensorflow 运行导出的模型等方式,可参见:Tensorflow 运行onnx
在移动端使用caffe2 可参考:移动端部署caffe2