Paddle模型转ONNX

深度学习模型在硬件加速器上的部署常常要用到ONNX(Open Neural Network Exchange,开放神经网络交换)格式,也可以通过ONNX实现不同AI框架(如Pytorch、TensorFlow、Caffe2、PaddlePaddle等)之间的模型转换。

Paddle模型转ONNX_第1张图片

最近临时使用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

你可能感兴趣的:(Python,深度学习,paddle)