深度学习模型在硬件加速器上的部署常常要用到ONNX(Open Neural Network Exchange,开放神经网络交换)格式,也可以通过ONNX实现不同AI框架(如Pytorch、TensorFlow、Caffe2、PaddlePaddle等)之间的模型转换。
最近临时使用PaddlePaddle开发模型,需要将准备好的模型转换成ONNX格式,在此记录一下。
Paddle转ONNX库需要用到paddle2onnx工具。在完成PaddlePaddle环境安装的基础上,通过如下命令安装paddle2onnx:
pip install paddle2onnx
安装完成后,就可以通过该工具完成paddle模型向onnx的转换了。下面就用一个简单的例子来说明用法。
首先,定义一个我们自己的模型。比如,我们定义一个轻量级OCR检测模型,该模型由三部分组成,分别是:一个轻量级特征提取网络MobileNetV3、一个金字塔结构FPN和一个Head,这三个结构已经在PaddleOCR开源代码中定义好了,我们只需要自己将这几部分拼成一个模型即可:
import paddle
from ppocr.modeling.heads.det_db_head import DBHead
from ppocr.modeling.backbones.det_mobilenet_v3 import MobileNetV3
from ppocr.modeling.necks.db_fpn import DBFPN
class MyPaddleOCR(paddle.nn.Layer):
def __init__(self, out_channels):
super(MyPaddleOCR, self).__init__()
self.backbone = MobileNetV3()
self.in_channels = self.backbone.out_channels
self.fpn = DBFPN(in_channels=self.in_channels, out_channels=out_channels)
self.db_head = DBHead(in_channels=out_channels)
def forward(self, x):
backbone_outs = self.backbone(x)
fpn_outs = self.fpn(backbone_outs)
head_outs = self.db_head(fpn_outs)
return head_outs
网络定义好之后,当然要进行模型训练,或者加载已知模型参数,这部分略去,我们直接讲如何将paddle模型转换成onnx。
模型转换需要用到paddle.onnx.export接口,像其他框架一样,导出模型时也需要指定输入数据的shape,在Paddle中,可以使用paddle.static.InputSpec来指定输入数据的shape,调用流程如下:
paddle_ocr_model = MyPaddleOCR(out_channels=256)
# ......
# 省略模型训练或参数加载
# 定义输出onnx文件路径
onnx_path = 'paddle_ocr_det'
# 定义输入数据,包括shape、type、name
# 将某一维度设置为None,则该维度为动态维度
#x_spec = paddle.static.InputSpec([None, 1, 640, 640], 'float32', 'x')
x_spec = paddle.static.InputSpec([1,3,640,640], 'float32', 'x')
# 调用paddle.onnx.export接口导出onnx模型
paddle.onnx.export(paddle_ocr_model, onnx_path, input_spec=[x_spec], opset_version=11)
好了,现在在当前目录下生成了paddle_ocr_det.onnx文件,让我们用onnxruntime来运行一下:
import onnxruntime
import numpy as np
onnx_path = 'paddle_ocr_det.onnx'
x = np.random.random((1, 3, 640, 640)).astype('float32')
ort_sess = onnxruntime.InferenceSession(onnx_path)
ort_inputs = {ort_sess.get_inputs()[0].name: x}
ort_outs = ort_sess.run(None, ort_inputs)
print(ort_outs[0].shape)
输出结果的形状为: (1, 3, 640, 640),是OK的。
ps:如果环境中未安装onnxruntime,需要先安装一下,看自己的机器是否有GPU,选择以下两个命令安装即可:
非GPU机器:
pip install onnxruntime
GPU机器:
pip install onnxruntime-gpu