从PyTorch、TensorFlow和PaddlePaddle模型转换为ONNX的完整教程

作者: memory

时间: 2024-10-18

在深度学习领域，模型转换是一个常见的需求，尤其是在不同框架之间进行迁移时。ONNX（Open Neural Network Exchange）作为一种开放的深度学习模型格式，允许我们在不同的框架之间轻松地转换模型。本文将详细介绍如何将PyTorch、TensorFlow和PaddlePaddle模型转换为ONNX格式，并提供详细的步骤和示例代码。

1. 准备工作

在开始之前，请确保你已经安装了以下工具和库：

Docker（用于创建一致的开发环境）
Python 3.x
PyTorch
TensorFlow
PaddlePaddle
ONNX
tf2onnx（用于TensorFlow模型转换）
paddle2onnx（用于PaddlePaddle模型转换）

你可以通过以下命令安装这些工具和库：

pip install torch tensorflow paddlepaddle onnx tf2onnx paddle2onnx

2. PyTorch模型转ONNX

2.1 创建工作目录

首先，创建一个工作目录来存放你的PyTorch模型和转换脚本。

mkdir torch_model
cd torch_model

2.2 搭建并保存模型

接下来，创建一个简单的PyTorch模型，并将其权重保存到文件中。创建一个名为simple_net.py的脚本，内容如下：

#!/usr/bin/env python3
import torch

# 构建一个简单的神经网络模型
class SimpleModel(torch.nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.m1 = torch.nn.Conv2d(3, 8, 3, 1, 0)
        self.m2 = torch.nn.Conv2d(8, 8, 3, 1, 1)

    def forward(self, x):
        y0 = self.m1(x)
        y1 = self.m2(y0)
        y2 = y0 + y1
        return y2

# 创建模型实例并保存其权重
model = SimpleModel()
torch.save(model.state_dict(), "weight.pth")

运行该脚本后，你将在当前目录下得到一个名为weight.pth的权重文件。

2.3 导出ONNX模型

现在，我们将PyTorch模型导出为ONNX格式。创建另一个名为export_onnx.py的脚本，内容如下：

#!/usr/bin/env python3
import torch
from simple_net import SimpleModel

# 加载预训练模型并导出为ONNX格式
model = SimpleModel()
model.eval()
checkpoint = torch.load("weight.pth", map_location="cpu")
model.load_state_dict(checkpoint)

# 准备输入张量
input = torch.randn(1, 3, 16, 16, requires_grad=True)

# 将PyTorch模型导出为ONNX格式
torch.onnx.export(model,
                  input,
                  'model.onnx',  # 导出的ONNX模型名称
                  opset_version=13,
                  export_params=True,
                  do_constant_folding=True)

运行该脚本后，你将在当前目录下得到一个名为model.onnx的ONNX模型文件。

3. TensorFlow模型转ONNX

3.1 创建工作目录

首先，创建一个工作目录来存放你的TensorFlow模型和转换脚本。

mkdir tf_model
cd tf_model

3.2 准备并转换模型

接下来，我们将使用TensorFlow官方提供的mobilenet_v1_0.25_224模型作为示例。通过以下命令下载模型并将其转换为ONNX格式：

wget -nc http://download.tensorflow.org/models/mobilenet_v1_2018_08_02/mobilenet_v1_0.25_224.tgz
tar xzf mobilenet_v1_0.25_224.tgz
python -m tf2onnx.convert --graphdef mobilenet_v1_0.25_224_frozen.pb 
    --output mnet_25.onnx --inputs input:0 
    --inputs-as-nchw input:0 
    --outputs MobilenetV1/Predictions/Reshape_1:0

运行以上命令后，你将在当前目录下得到一个名为mnet_25.onnx的ONNX模型文件。

4. PaddlePaddle模型转ONNX

4.1 创建工作目录

首先，创建一个工作目录来存放你的PaddlePaddle模型和转换脚本。

mkdir pp_model
cd pp_model

4.2 准备模型

接下来，我们将使用PaddlePaddle官方提供的SqueezeNet1_1模型作为示例。通过以下命令下载模型并进行形状推理：

wget https://bj.bcebos.com/paddlehub/fastdeploy/SqueezeNet1_1_infer.tgz
tar xzf SqueezeNet1_1_infer.tgz
cd SqueezeNet1_1_infer

下载并运行paddle_infer_shape.py脚本，设置输入形状为[1, 3, 224, 224]：

wget https://raw.githubusercontent.com/PaddlePaddle/Paddle2ONNX/develop/tools/paddle/paddle_infer_shape.py
python paddle_infer_shape.py  --model_dir . 
                          --model_filename inference.pdmodel 
                          --params_filename inference.pdiparams 
                          --save_dir new_model 
                          --input_shape_dict="{'inputs':[1,3,224,224]}"

运行完以上命令后，你将在当前目录下得到一个名为new_model的目录。

4.3 转换模型

最后，使用paddle2onnx工具将PaddlePaddle模型转换为ONNX格式：

pip install paddle2onnx
paddle2onnx  --model_dir new_model 
          --model_filename inference.pdmodel 
          --params_filename inference.pdiparams 
          --opset_version 13 
          --save_file squeezenet1_1.onnx

运行完以上命令后，你将得到一个名为squeezenet1_1.onnx的ONNX模型文件。

如果你在转换过程中遇到任何问题，可以参考ONNX官方仓库提供的教程：ONNX Tutorials。