如何将PyTorch模型从ckpt转换为ONNX格式？

Question

**如何将PyTorch模型从ckpt文件转换为ONNX格式？** 在深度学习模型部署过程中，常需将训练好的PyTorch模型（通常保存为`.ckpt`或`.pt`格式）转换为ONNX格式以实现跨平台兼容性。然而，许多开发者在加载检查点文件、重建模型结构及执行导出时遇到困难。本文将详细介绍如何正确加载PyTorch模型、构造输入张量，并使用`torch.onnx.export`接口将其转换为ONNX格式，涵盖常见问题如模型结构不匹配、输入维度错误等，帮助你顺利完成模型转换与部署。

Answer 1

一、PyTorch模型与ONNX格式简介

在深度学习模型部署过程中，常常需要将训练好的模型转换为通用的中间表示格式，以便于在不同平台和框架中进行推理。ONNX（Open Neural Network Exchange）是一种开放的神经网络交换格式，支持多种深度学习框架之间的模型互操作。

PyTorch通常使用`.pt`或`.ckpt`文件保存模型状态字典（state_dict），这些文件仅包含模型参数，并不包含完整的模型结构定义。因此，在将模型转换为ONNX格式之前，必须先重建模型结构并加载对应的参数。

• ONNX的优势：跨平台兼容、支持多种推理引擎（如TensorRT、ONNX Runtime等）
• PyTorch模型保存方式：state_dict / 整体模型保存

二、准备环境与依赖项

为了顺利完成从`.ckpt`到ONNX的转换，需确保以下软件包已安装：

库名	版本要求
torch	>=1.8.0
onnx	>=1.9.0
onnxruntime	>=1.8.0

可通过如下命令安装必要依赖：

pip install torch onnx onnxruntime

三、加载PyTorch模型检查点

PyTorch模型通常以两种方式保存：

1. 仅保存模型参数（state_dict）
2. 保存整个模型（model.save）

若使用的是state_dict方式保存的`.ckpt`文件，则需手动重新构建模型结构后再加载参数。

# 示例：加载state_dict
import torch
from model_definition import MyModel  # 假设这是你的模型定义

model = MyModel()
checkpoint = torch.load('model.ckpt')
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()  # 设置为评估模式

四、构造输入张量

导出ONNX模型时，必须提供一个或多个示例输入张量，用于追踪模型执行路径。

输入张量的维度应与训练/推理时一致。例如，图像分类任务中常见的输入形状为 (batch_size, channels, height, width)。

# 构造虚拟输入
dummy_input = torch.randn(1, 3, 224, 224)  # batch_size=1, 3通道，224x224图像

注意：如果模型有多个输入，可以传入元组形式。

五、使用torch.onnx.export接口导出ONNX模型

PyTorch提供了`torch.onnx.export`函数用于导出模型至ONNX格式。

import torch.onnx

# 导出ONNX模型
torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    export_params=True,  # 存储训练参数
    opset_version=13,    # ONNX算子集版本
    do_constant_folding=True,  # 优化常量
    input_names=['input'],     # 输入节点名称
    output_names=['output'],   # 输出节点名称
    dynamic_axes={
        'input': {0: 'batch_size'},  # 动态维度
        'output': {0: 'batch_size'}
    }
)

该函数支持多种参数配置，可满足静态图和动态图的需求。

六、验证ONNX模型正确性

导出后，建议使用ONNX Runtime进行推理测试，以确保模型输出与原始PyTorch模型一致。

import onnx
import onnxruntime as ort
import numpy as np

# 加载ONNX模型
onnx_model = onnx.load("model.onnx")
onnx.checker.check_model(onnx_model)

# 创建推理会话
ort_session = ort.InferenceSession("model.onnx")

# 运行推理
outputs = ort_session.run(
    None,
    {'input': dummy_input.numpy()}
)
print(outputs[0])

七、常见问题与解决方案

1. 模型结构不匹配：确认模型定义与保存的state_dict完全一致。
2. 输入维度错误：确保dummy_input与训练时输入一致，尤其注意通道顺序。
3. 导出失败，提示未支持的操作：尝试更新PyTorch版本或使用更高级的opset版本。
4. 动态轴未生效：检查dynamic_axes参数是否正确定义。

八、进阶技巧与优化建议

• 使用ONNX TensorRT加速推理
• 利用ONNX Runtime进行多平台部署
• 使用自定义ONNX符号化函数支持特定层

graph TD A[开始] --> B{是否有完整模型结构?} B -- 是 --> C[直接加载模型] B -- 否 --> D[手动定义模型结构] D --> E[加载state_dict] C --> E E --> F[构建输入张量] F --> G[调用torch.onnx.export] G --> H[生成ONNX模型] H --> I[使用ONNX Runtime验证]