Paddle模型转ONNX

news2024/11/17 9:25:31

深度学习模型在硬件加速器上的部署常常要用到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机器：