【人工智能技术应用学习笔记09】PyTorch 模型如何转换成 ONNX- 惊觉

【人工智能技术应用学习笔记09】PyTorch 模型如何转换成 ONNX

一句话导读：上一篇认识了 ONNX 是什么，这一篇动手把一个 PyTorch 模型真正变成 .onnx 文件，并用 ONNX Runtime 验证转换结果。

系列：《人工智能技术应用学习笔记》
第 9 篇：PyTorch 模型如何转换成 ONNX
关键词：PyTorch、ONNX、torch.onnx.export、onnxruntime、模型部署

前言

上一篇《ONNX 是什么？》里我们说到：ONNX 是一种中立的深度学习模型表示规范，它把 PyTorch、TensorFlow 训练出来的模型统一成一种格式，方便交给 TensorRT、ONNX Runtime 等工具去部署。

但"知道 ONNX 是什么"只是第一步。真正落地时，我们要能把自己训练好的 PyTorch 模型，准确无误地导出一个 .onnx 文件。这篇博客就来做这件事。

本次目标

掌握 torch.onnx.export 的基本用法；
把一个 PyTorch 模型导出为 .onnx 文件；
用 onnx 库检查导出的模型是否合法；
用 onnxruntime 跑一遍推理，并与 PyTorch 的输出对比，证明转换正确。

环境说明

项目	版本
操作系统	Windows 10 / 11
Python	3.10
PyTorch	2.2.1（GPU / CPU 均可）
onnx	1.16.0
onnxruntime	1.17.1

说明：torch.onnx.export 是 PyTorch 自带的，不需要额外安装，PyTorch ≥ 1.8 都可以使用。onnx 和 onnxruntime 是后面做验证时用的，需要另外安装。

原理简介

torch.onnx.export 的核心原理是追踪（Tracing）：

给模型喂一个样例输入（dummy input）；
PyTorch 拿着这个输入，把模型 forward 从头到尾真实执行一遍；
把执行过程中经过的每一个算子（如 Conv、ReLU、Linear）记录下来，整理成一张静态计算图；
再按照 ONNX 规范，把这张图序列化成 .onnx 文件。

因为是"跑一遍再记录"，所以有两点要特别注意：

控制流：如果模型里存在依赖输入值的 Python 逻辑（比如 if x > 0: 这种），追踪时只会记录当时走的那条分支，可能导致导出错误，后面常见问题里会细说；
动态维度：导出时 batch 大小会默认被"固定"成样例输入的大小，如果希望部署时能接受任意 batch，需要显式声明 dynamic_axes。

另外还有两个常用参数：

opset_version：ONNX 算子的版本集合，数字越大算子越新。一般选 13～17 之间，同时要和部署端（TensorRT / ONNX Runtime）兼容；
input_names / output_names：给输入输出起名字，方便部署端按名字调用。

操作过程

1. 安装验证依赖

pip install onnx onnxruntime

2. 定义一个简单的模型

为了不让示例依赖太多，先用一个最简单的三层网络做演示：

import torch
import torch.nn as nn

class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(4, 8),
            nn.ReLU(),
            nn.Linear(8, 2),
        )

    def forward(self, x):
        return self.net(x)

model = SimpleNet()
model.eval()  # 关键：切换到推理模式，关闭 Dropout / BatchNorm 的训练行为

3. 准备一个样例输入

导出必须提供一个"形状正确"的样例输入，PyTorch 会用它来跑一遍 forward：

dummy_input = torch.randn(1, 4)  # batch=1，特征维度=4

4. 调用 torch.onnx.export 导出

torch.onnx.export(
    model,                      # 要导出的模型
    dummy_input,                # 样例输入
    "simple_net.onnx",          # 输出文件名
    input_names=["input"],      # 输入节点名字
    output_names=["output"],    # 输出节点名字
    dynamic_axes={              # 声明 batch 维度为动态
        "input": {0: "batch_size"},
        "output": {0: "batch_size"},
    },
    opset_version=17,           # 算子集版本
)
print("导出完成：simple_net.onnx")

常用参数说明：

参数	作用
`model`	要导出的 PyTorch 模型
`args`	样例输入，可以是单个 Tensor 或元组
`f`	输出的 .onnx 文件路径
`input_names` / `output_names`	输入/输出节点的名字
`dynamic_axes`	声明哪些维度是动态的（如 batch）
`opset_version`	使用的 ONNX 算子集版本

验证结果

导出后不要直接拿去用，先做两件验证，确认转换没问题。

1. 用 onnx 检查模型是否合法

import onnx

onnx_model = onnx.load("simple_net.onnx")
onnx.checker.check_model(onnx_model)  # 不报错就说明模型结构合法

print("模型检查通过，算子数量：", len(onnx_model.graph.node))

2. 用 onnxruntime 推理并和 PyTorch 对比

import numpy as np
import onnxruntime as ort

# 创建 ONNX Runtime 会话
ort_session = ort.InferenceSession("simple_net.onnx", providers=["CPUExecutionProvider"])

# 用同一份输入跑 ONNX
ort_inputs = {"input": dummy_input.numpy()}
ort_output = ort_session.run(None, ort_inputs)[0]

# 用 PyTorch 跑一遍做对照
with torch.no_grad():
    torch_output = model(dummy_input).numpy()

print("PyTorch 输出：", torch_output)
print("ONNX 输出：", ort_output)
print("最大误差：", np.max(np.abs(torch_output - ort_output)))

运行结果类似：

PyTorch 输出： [[-0.3215  0.0893]]
ONNX 输出： [[-0.3215  0.0893]]
最大误差： 1.1920929e-07

只要"最大误差"是一个接近 0 的小数（如 1e-7 量级），就说明转换后的 ONNX 模型和原 PyTorch 模型行为一致，可以放心拿去部署了。

补充：如果模型里没有随机性操作，并且正确调用了 model.eval()，两边结果在 float32 下通常只差极小的浮点误差，属于正常现象。

常见问题

导出的 batch 被固定成 1，部署时换 batch 就报错？
导出时用 dynamic_axes 把 batch 维度声明成动态的，如上面的 {0: "batch_size"}。
模型里写了 if / for 等依赖输入的控制流，导出后结果不对？
因为 Tracing 只记录实际执行到的那条路径。解决办法：改用脚本化 torch.jit.script 重写模型，或者把控制流改成 Tensor 运算（如 torch.where）。
报错提示某个算子不支持 / opset 版本问题？
尝试降低 opset_version（比如从 17 降到 13），或者升级 PyTorch / onnx 版本。注意 opset 要和部署端兼容。
自定义算子 / 很新的算子导不出来？
ONNX 标准里没有的算子需要自己注册。日常学习阶段，尽量用常见的标准算子（Conv、Linear、ReLU 等）就能避免这个问题。
导出后精度差异比较大？
先检查是否忘了 model.eval()，以及模型里是否有 Dropout、BatchNorm 这类在训练/推理模式行为不同的层。同时建议用和导出时一致的 float32 精度做对比。

总结

本篇把"PyTorch 模型 → ONNX"这一步完整走了一遍：

核心函数只有一个：torch.onnx.export；
导出前要 model.eval() 并准备一个样例输入；
用 dynamic_axes 声明动态维度，用 opset_version 控制算子版本；
导出后要用 onnx.checker 检查合法性，再用 onnxruntime 跑一遍和 PyTorch 对比，误差接近 0 才算成功。

到这里，模型已经从"训练框架"走到了"通用中间格式"。下一步就是把这个 ONNX 模型交给部署工具做加速。

下一篇预告

《人工智能技术应用学习笔记10：ONNX 模型如何转换为 TensorRT Engine》

下一篇会继续讲：

TensorRT Engine 是什么；
怎么把 .onnx 文件转换成 TensorRT 的 engine；
转换时 FP32 / FP16 / INT8 该怎么选。

标签：#人工智能 #深度学习 #PyTorch #ONNX #学习 #笔记

转载自 CSDN-专业IT技术社区

原文链接：https://blog.csdn.net/xinjiyaya/article/details/168242936

【人工智能技术应用学习笔记09】PyTorch 模型如何转换成 ONNX