PyTorch模型转ONNX全流程指南：从原理到部署避坑

你训练完一个 PyTorch 模型，第一反应通常是“精度到了，搞定收工”。但真正的问题是：这个模型去哪儿跑？它要跑在树莓派上，还是接入公司的 C++ 服务？还是塞进瑞芯微开发板做 int8 量化？PyTorch 模型本身不方便直接搬到这些环境里，所以在模型部署这条路上，第一步几乎绕不开一个词：onnx。

我用 PyTorch 也做了好几年模型训练和部署，最初被“模型部署”这事折腾得够呛。后来发现，只要把 PyTorch 模型正确转成 onnx，后面的推理、量化、跨平台移植都会顺很多。这篇文章从最实操的角度，把“PyTorch 模型转 onnx”的来龙去脉、关键参数、坑点和验证手段完整梳理一遍，适合刚接触模型部署的算法工程师，也适合要把自己训练过的 YOLOv5 之类模型往边缘设备上搬的应用开发者。

1. 模型转换的核心逻辑：为什么要让 PyTorch 模型变成 onnx

1.1 部署场景天然要求中间格式

PyTorch 模型本身不是一种“部署友好”的格式。训练用的 model.state_dict() 只是权重字典，加载它时需要重建完整网络结构，还依赖 PyTorch 版本、自定义层代码，甚至 Python 运行环境。如果把训练好的模型直接交给部署同学，对方要装 PyTorch、要装 CUDA 对应版本、还要保证网络代码不被改动，这在真实业务里极容易翻车。

所以业界的常见做法是，在训练框架和推理框架中间插一个中间格式，让训练框架负责导出，让推理框架负责加载执行。onnx（Open Neural Network Exchange）就是这个中间层的核心代表，它定义了一套计算图规范，把网络结构、算子、权重统一保存下来，和 PyTorch 不再是强绑定关系。导出的 onnx 文件本质上是一个带权重的计算图，推理端拿到它之后，不需要 PyTorch 也能跑。

说得更直白一点：onnx 之于模型部署，就像通用图片格式之于文档发布。你总不能要求每个看文档的人都装一份微软 Office，同样，你不能要求每个跑模型的服务都塞一个 PyTorch 环境。

1.2 onnx 在部署链路中的真实位置

onnx 不是终点，它是中转站。模型转成 onnx 之后，通常还会流向几个方向：

• 直接交给 onnxruntime 推理引擎，在 CPU/GPU 上跑，最简单、最通用。
• 转成 OpenVINO 中间表示，在 Intel CPU/核显上优化运行。
• 转成 TensorRT engine，在 NVIDIA GPU 上获得极致推理性能。
• 转成 RKNN 等硬件平台格式，在瑞芯微等边缘 NPU 上做 int8 量化推理。
• 用 onnx 自带的量化工具做 int8 量化后，再部署到移动端或嵌入式设备。

很多热搜场景，比如“瑞芯微转换onnx模型”“树莓派5上部署自己训练的yolov5模型”，实际路径都是 PyTorch -> onnx -> 目标格式。瑞芯微的 RKNN Toolkit 官方推荐先导出 onnx 再转 rknn；树莓派上也可以用 onnxruntime 直接跑你导出的 onnx 文件。

从我的经验看，onnx 最大的价值不在于它本身跑得有多快，而是它让“模型真正脱离训练环境”这件事变成现实。你的模型只要导成 onnx，后续无论是性能优化、框架更换、还是硬件迁移，都有一个可靠起点。

1.3 什么样的模型适合转 onnx

绝大多数 CV 模型和 NLP 模型都能顺利转 onnx，但也有一些例外。如果网络结构里使用了非常新的算子、依赖于某些 PyTorch 内部特殊实现，或者包含了无法静态 trace 的控制流，导出时会遇到算子不支持、转换失败等问题。

因此，建议在建模阶段就考虑部署约束。训练时尽量使用常见算子，比如 Conv、BatchNorm、ReLU、Linear、LayerNorm 等；少用 torch.where 导致大范围数据依赖的复杂分支，少依赖 Python 级别的 for 循环控制流。如果非要用，也可以，但需要在导出阶段做额外处理，后面我会讲。

2. 环境准备：版本匹配是转换成功的一半

2.1 PyTorch、onnx、onnxruntime 的安装顺序

先说结论：环境安装不复杂，但版本搭配很关键。我在实际项目里踩过太多次“onnx 报错，最后发现是 torch 和 onnx 版本不匹配”的坑。

建议在一个干净的 Python 虚拟环境里安装，避免把训练环境和部署相关依赖混在一起。我习惯用 conda 先建一个独立环境，比如部署专用环境叫 deploy：

bash复制conda create -n deploy python=3.10 -y
conda activate deploy
pip install torch torchvision --index-url https://download.pytorch.org/whl/cpu
pip install onnx onnxruntime

如果你本机是 NVIDIA GPU 环境，安装带 CUDA 版本的 PyTorch 也没问题，因为导出 onnx 时是否用 GPU 关系不大，核心计算发生在 CPU 上做 graph trace 和权重序列化。但 onnxruntime 需要根据目标环境选择 CPU 版还是 GPU 版，树莓派、嵌入式 Linux、mac 本地用 CPU 版即可。

2.2 版本对照参考

onnx 的算子版本（opset）与 PyTorch 版本关系密切。我用过不少版本组合，下面这张表是经过实际项目验证的相对稳定的搭配，供你参考：

安装后可以快速验证核心依赖是否就绪：

bash复制python -c "import torch; print(torch.__version__)"
python -c "import onnx; print(onnx.__version__)"
python -c "import onnxruntime; print(onnxruntime.__version__)"

如果 onnxruntime 安装不上，可以单独指定源安装，比如在某些内网环境下选择离线 wheel。Ubuntu 上如果系统 Python 环境被系统包管理器托管，强烈建议用虚拟环境，避免把系统搞乱。

2.3 树莓派和边缘设备的环境注意点

如果你打算把模型放到树莓派 5 上跑，可以预先在树莓派上安装 onnxruntime：

bash复制pip install onnxruntime

树莓派是 ARM 架构，onnxruntime 官方发布版本里有对应的 manylinux aarch64 wheel，直接安装即可。不过要注意，树莓派的 CPU 算力有限，大模型即使转成 onnx 也不一定跑得动。我自己在树莓派 5 上跑过 YOLOv5s 的 onnx 模型，CPU 推理一帧大约几百毫秒，能接受但谈不上流畅。如果后续要上 NPU、RKNN 这类硬件加速，那就必须把 onnx 再转成对应硬件格式，并且通常要求静态输入尺寸。

3. 导出前的模型准备：别上来就写 export

3.1 重新加载模型再导出，不要直接导出训练内存里的实例

不少初学者在训练脚本末尾接着导出当前模型，但这种做法很容易把训练状态混进去。模型在训练过程中可能处于 train() 模式，BatchNorm 层还在更新均值方差；也可能包含梯度缓存。导出的计算图可能带着脏状态。

正确做法是，从保存好的 checkpoint 中重新构造模型，先加载权重，再切换成 eval() 模式，最后做导出。示例：

python复制import torch
from models.detector import Detector

model = Detector(num_classes=20)
checkpoint = torch.load("best.pt", map_location="cpu")
model.load_state_dict(checkpoint["model_state_dict"])
model.eval()

这样能确保导出的模型结构和训练时一致，同时权重也是最终收敛版本。

3.2 eval 模式与开关梯度的影响

PyTorch 模型默认是训练模式，BatchNorm 层会使用当前 batch 的统计量，Dropout 层会随机失活。导出 onnx 时必须切换成 eval() 模式，把 BatchNorm 切换成使用训练阶段累计的 running_mean/running_var，Dropout 变成恒等映射。

同时，导出前包一段 with torch.no_grad(): 也是好习惯。虽然 torch.onnx.export 内部会做 gradient 隔离，但显式无梯度上下文可以让代码意图更清晰，避免意外触发某些自定义 autograd 函数。

python复制with torch.no_grad():
    torch.onnx.export(
        model,
        dummy_input,
        "model.onnx",
        input_names=["images"],
        output_names=["output"],
        opset_version=17,
        dynamic_axes=None,
    )

3.3 构造合适的 dummy input

导出 onnx 需要给模型一个真实的输入张量，这个张量不参与实际预测，只用于 trace 网络的张量流动，所以一般叫 dummy_input。它的大小要和实际部署尺寸匹配。

常见错误有两种：一种是直接用随机初始化张量，比如 torch.randn(1, 3, 640, 640)，这没问题；另一种是用全 1 张量，虽然也能 trace，但如果网络里有一些和输入分布强相关的统计操作，全 1 可能触发异常路径，所以还是推荐用随机数。

如果你的模型训练时输入尺寸是 640x640，那就用 (1, 3, 640, 640)；如果你的部署场景可能输入不同尺寸图片，那要把 dynamic_axes 配好，后面会讲。dummy input 的通道顺序也必须是模型实际需要的格式，常见 CV 模型都是 NCHW，也就是 batch、channel、height、width。

4. 实操核心：用 torch.onnx.export 把模型完整导出

4.1 最基础的可运行导出代码

直接给出一份我常用的、经过多个项目验证的导出脚本骨架，以常见目标检测模型为例：

python复制import torch
import onnx
import onnxruntime as ort
import numpy as np

model = Detector(num_classes=20)
checkpoint = torch.load("best.pt", map_location="cpu")
model.load_state_dict(checkpoint["model_state_dict"])
model.eval()

dummy_input = torch.randn(1, 3, 640, 640, device="cpu")

torch.onnx.export(
    model,
    dummy_input,
    "detector.onnx",
    export_params=True,
    opset_version=17,
    do_constant_folding=True,
    input_names=["images"],
    output_names=["boxes", "scores", "labels"],
    dynamic_axes=None,
)

onnx_model = onnx.load("detector.onnx")
onnx.checker.check_model(onnx_model)
print("onnx model converted successfully")

这个脚本短短几十行，已经把最关键的工作做完了。但要真正理解里面每个参数的含义，否则后续遇到奇葩问题会一头雾水。

4.2 参数选择：不看文档也能用对的五个关键点

torch.onnx.export 常用参数不少，但决定成败的核心参数就下面几个：

• export_params=True：是否把模型权重保存到 onnx 文件内。部署时绝对要设为 True，否则导出的 onnx 只有计算图没有权重，拿给别人根本加载不了。
• opset_version：onnx 算子版本。不要一味追求高版本，如果部署环境是老版本 onnxruntime，高版本算子可能不支持。建议低版本环境用 11，通用场景用 13 或 17。RKNN Toolkit 等工具通常支持到 11/12，导出前最好先确认硬件厂商要求的 opset 范围。
• do_constant_folding=True：让导出工具自动折叠常量计算，减少图里的冗余节点，一般打开能减小模型体积、加快推理。但如果你之后要做 int8 量化或模型结构修改，有时量化工具会抱怨折叠后的图结构不直观，这种情况下可以考虑关掉再导出。
• input_names 和 output_names：给模型的输入输出节点起名字。起一个好记的名字很重要，因为后续转 RKNN、OpenVINO、TensorRT 时，你经常会用到这些名字来指定输入输出。例如 RKNN Toolkit 转换时就需要指定输入节点名称。
• dynamic_axes：指定哪些维度是动态的。这个是部署中最容易出问题的点，单独拿出来说。

4.3 动态轴配置详解：一张图上不同尺寸输入的方案

如果你的模型在实际部署时要接收任意宽高图片，比如实时视频流，那么输入尺寸不能固定在 640x640。此时需要配置 dynamic_axes，告诉 onnx 哪些维度允许变化：

python复制dynamic_axes = {
    "images": {0: "batch_size", 2: "height", 3: "width"},
    "boxes": {0: "num_boxes"},
    "scores": {0: "num_boxes"},
    "labels": {0: "num_boxes"},
}

torch.onnx.export(
    model,
    dummy_input,
    "detector_dynamic.onnx",
    opset_version=17,
    input_names=["images"],
    output_names=["boxes", "scores", "labels"],
    dynamic_axes=dynamic_axes,
)

配置动态轴后，导出的 onnx 模型在推理时允许 batch 大小变化、图像宽高变化。但有几个注意点：

第一，动态轴会降低部分推理框架的优化力度。TensorRT、RKNN 这类工具对动态输入往往需要额外的优化配置，甚至不支持，所以很多边缘部署场景干脆把分辨率固定，用静态图换取最大性能。

第二，onnxruntime 虽然支持动态轴，但如果输入图像尺寸变化范围过大，内部会反复重新分配内存，实际吞吐不一定好看。

第三，动态轴只对维度生效，如果你的模型内部有“对图像大小做 reszie 后固定到某个尺寸”的逻辑，那外部输入再动态，网络核心部分也是静态的，收益有限。

从部署稳定性的角度，我个人的建议是：如果目标平台是服务器 GPU，开动态轴问题不大；如果目标是 RKNN、树莓派或者其他边缘 NPU，优先保证静态尺寸导出，先跑通再谈动态。

4.4 opset 版本的选择逻辑

onnx 每代版本会增加一些算子或修改算子定义，PyTorch 导出的算子会根据 opset_version 选择。比如老版本 opset 不支持 Resize 算子的某些属性，要用 Upsample 替代，而新版本又倾向于 Resize。

如果你的部署链路上有 RKNN 工具链，务必查一下他们官方支持的最高 opset。有些硬件平台很老，只支持 opset 11，你导出成 17，转换工具会直接报算子版本不兼容。反过来，如果部署目标是 ONNX Runtime 最新版，那 opset 17 是比较稳妥的通用选择。

我平时默认导出两版：一版是 opset 11，用于兼容老工具链；另一版是 opset 17，用于现代 onnxruntime 推理。实际部署时根据目标环境“只选对的不选高的”。

5. 验证：导出 onnx 只是开始，跑通并对比精度才是关键

5.1 onnx.checker 只能查“结构病”

很多教程在导出后做个 checker 就完事了，但 onnx.checker.check_model() 检查的是图结构是否合法，例如节点输入输出能否对应上、属性是否完整。它不会告诉你模型推理结果对不对，更不能验证精度损失。

所以更重要的验证步骤是，用 onnxruntime 加载导出的模型，给它相同的输入，对比 PyTorch 模型输出和 onnxruntime 输出是否一致。

5.2 PyTorch 推理和 onnxruntime 推理对比验证

下面这段代码我基本每次导出后都会跑一遍，可以作为你的验证模板：

python复制import onnxruntime as ort
import numpy as np
import torch

def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()

test_input = torch.randn(1, 3, 640, 640)
with torch.no_grad():
    torch_outputs = model(test_input)

sess = ort.InferenceSession("detector.onnx", providers=["CPUExecutionProvider"])
ort_inputs = {sess.get_inputs()[0].name: to_numpy(test_input)}
ort_outputs = sess.run(None, ort_inputs)

for i, (torch_out, ort_out) in enumerate(zip(torch_outputs, ort_outputs)):
    np.testing.assert_allclose(to_numpy(torch_out), ort_out, rtol=1e-3, atol=1e-5)
    print(f"output {i} shape: {ort_out.shape}, max diff: {np.max(np.abs(to_numpy(torch_out) - ort_out))}")

最终输出结果可能并不是逐元素完全一致，因为部分算子在不同框架里的浮点实现存在微小差异，但一般要求误差在 1e-3 到 1e-4 级别。如果差异到了 1e-1 甚至更大，那基本可以断定导出的图有问题。

比较时要注意：如果你的模型前处理包含归一化或其他前置步骤，请先手动完成同样的预处理再接 onnx 推理，否则结果天然会错开。很多次“精度对不齐”根本不是转换问题，而是输入数据的处理流程没对齐。

5.3 推理性能验证

除了精度，还要看一眼转换后的 onnx 在 onnxruntime 下的推理耗时。这部分数据对你评估部署性能非常重要。

python复制import time

# warm up
for _ in range(10):
    sess.run(None, ort_inputs)

start = time.time()
count = 100
for _ in range(count):
    sess.run(None, ort_inputs)
cost = (time.time() - start) / count * 1000
print(f"average inference time: {cost:.2f} ms")

注意，树莓派等 CPU 平台首次推理可能包含算子初始化开销，统计前先 warm up 是必要的。真实场景里的耗时还受线程数影响，onnxruntime 默认会尽量使用多核，但有些边缘设备希望限制 CPU 占用，可以通过 sess_options.intra_op_num_threads 控制。

5.4 与精度量化相关的部署考量

如果后续要做 int8 量化，那么这一阶段还需要额外关注 onnx 图里每个节点的数值范围。RKNN 或 onnxruntime 的 int8 量化通常需要校准数据，而校准数据应当尽量贴近真实部署输入分布。你现在导出的 fp32 onnx 就是量化的输入，必须在导出阶段保证算子表达正确，否则后面量化后的精度损失可能会被误以为是量化本身的问题。

6. 导出过程中的常见坑与排查方法

6.1 常见错误对照表

下面是我在实际项目中频繁遇到、且经常让新同事卡住的几类问题，整理成速查表：

6.2 算子不支持时的处理思路

如果遇到 Unsupported operator，第一步是升级 opset 版本再看。如果升级后仍不支持，说明 PyTorch 导出器没有把该算子注册到 onnx 算子集里，这时有两条路可以走。

一条是从网络结构上规避，把自定义模块改写成常见算子组合。比如有些实现喜欢用 torch.repeat_interleave，它在某些旧版本 opset 下不支持，可以改成用 repeat + reshape 组合实现同样效果，既符合 onnx 语义，又方便后续硬件移植。

另一条是为自定义算子写符号注册函数，相当于告诉 PyTorch 如何把你自己实现的 forward 翻译成 onnx 算子。这种方案适合团队内部有公共自研层的情况，但对新手来说门槛偏高，优先推荐改造网络结构。

如果用的是 YOLOv5 这类成熟检测模型，基本不需要碰自定义算子。官方仓库已经内置了导出脚本，只需要调好 dynamic 参数就能导出自己的权重。

6.3 输入输出名称与后续工具链的配合

很多情况下，导出的 onnx 能在 onnxruntime 跑通，但拿到瑞芯微或 OpenVINO 工具链里就报找不到输入或输出，原因多半是节点名称对不上。不同工具链对输入输出的要求和显示名不一样，所以在导出时就要规划好命名。

我建议固定一套命名规范：图像输入统一叫 images，文本输入统一叫 input_ids，检测输出统一用 boxes、scores、labels。如果模型只有单个输出张量，则建议叫 output。这样后续无论转 RKNN 还是写推理 C++ 代码，看到的都是熟悉的名字，不会每次都被工具链的报错牵着跑。

如果你拿到的 onnx 模型是别人导出的、命名不友好，可以用 onnx.utils.extract_model 提取子图并重新指定输入输出名，不过能做净化的前提是图结构本身清晰。

6.4 与模型输入预处理相关的隐蔽坑

导出的 onnx 模型通常不包含图像解码和归一化，除非你在模型 forward 里手动接入了这些步骤。因此，在 PyTorch 里推理时如果对图像做了 mean/std 归一化、BGR/RGB 通道交换，那么在 onnxruntime 里跑同一张图时也必须做同样操作，否则结果差异会非常明显。

最容易翻车的一个点是 PyTorch 中图像通道顺序是 [batch, channel, height, width]，但很多常见图像库以 HWC 形式输出。如果你在导出前没有固定输入张量格式，部署端又搞混了 NCHW 和 NHWC，模型推理出来的结果大概率是乱的。解决方法是做一个简单的单元测试：让你的预处理输出经过 PyTorch 推理得到 A，再让同一份输入经过 onnxruntime 推理得到 B，比对是否一致。

6.5 模型内部控制流和动态形状的报错

PyTorch 里 if 条件如果依赖输入张量的实际值，而不是形状等静态信息，torch.onnx.export 使用的 trace 机制只能记录那次输入下执行的某一分支，无法覆盖另一端条件。这样导出的模型实际上是“残缺”的。

应对办法是，尽量把动态条件变成静态条件。比如，把“如果 batch size 为 1 就走单样本逻辑，否则走 batch 逻辑”这类代码统一成一个批量路径。或者，如果必须保留运行时分支，可以考虑设计成多个不同输入尺寸的静态分支，或者使用 onnxruntime 本身不直接支持的动态控制流——这种情况通常要单独讨论，不适合普通部署。

7. 模型导出后的延伸部署思路

7.1 不同目标平台下的下一步动作

onnx 转换完成后，真正的部署可以根据目标平台选不同路径。下面是我常见的几条分支：

如果目标平台是树莓派这类通用 ARM 设备，onnxruntime 直接跑就够用了，贪图再转一套 NCNN 未必获得明显收益，反而增加转换工作量。

7.2 静态图与动态图的选择原则

模型从 PyTorch 转成 onnx，图形是静态的还是动态的，会极大影响后面能做的事情。静态图意味着输入尺寸、batch 大小都是固定的，这样推理框架可以在执行前做大幅图优化，硬件的内存分配也可以预先规划。动态图则灵活，但很多算子级优化无法提前做，运行时会多出很多 shape 推断开销。

从实际部署看，边缘硬件平台往往先用静态图，比如树莓派上的 YOLOv5 如果固定为 640x640 输入，跑起来会比动态尺寸稳定不少。服务端如果请求的图片尺寸乱七八糟，才优先考虑动态轴，避免每次重新缩放图片。

我个人的做法是“能静态就静态，必须动态再开动态”。当你的模型不止在一种场景使用时，建议为不同场景各导出一版专用 onnx，而不是让一个动态模型包打天下。

7.3 onnx 模型的可视化与网络结构检查

拿到导出的 onnx 文件后，想看内部结构，推荐用 Netron 打开，直接在浏览器里查看网络节点、输入输出、参数 shape，非常直观。它是调试部署问题的效率工具，建议每台搞推理的电脑上都装一个。

另外，onnx 自带 onnx.shape_inference.infer_shapes，可以用来推断 onnx 各节点的中间张量形状。如果结构里出现一些导出的节点输出维度为 None，你至少能提前知道动态 shape 的影响范围。

python复制import onnx
from onnx import shape_inference

model = onnx.load("detector.onnx")
model = shape_inference.infer_shapes(model)
onnx.save(model, "detector_inferred.onnx")

这种 infer 后的模型有时候在部署阶段可以减少推理时的动态 shape 推断工作量。

7.4 从 onnx 到 RKNN 的额外注意事项

因为热搜词里多次出现瑞芯微，这里单独补充几句。瑞芯微 NPU 的官方工具链 RKNN-Toolkit2 通常要求输入模型是 onnx 或 pytorch，但转换成 rknn 时，它只能支持有限的 opset 范围，而且对动态形状支持较差。

所以在为 RKNN 准备模型时，导出 onnx 阶段注意几条：

• 确定好最终推理分辨率，固定宽高导出。
• 确定好颜色通道顺序，尽量不要在网络里做 RGB/BGR 切换。
• 如果模型里包含一些自定义的后处理（例如 NMS 后处理在 PyTorch 侧实现），通常建议先导出不带后处理的裸网络输出，再用 NPU 或 CPU 实现后处理。
• 算子越简单越好。如果 onnxruntime 都支持且没问题，RKNN 工具还不支持，往往需要回到 PyTorch 代码里替换算子重新导出。

写在最后的几个实际体会

模型部署这件事，看起来就是一行 export 的事，但水比想象中深。我见过太多人在这一步卡住，其实多数问题不是玄学，而是对 export 的机制、算子集版本和平台要求没有理解透。

真正高效的做法是，把“导出 onnx”当成一条流水线，而不是单个动作：从评估模型结构、准备环境、构造 dummy input、设置导出参数，到用 onnxruntime 反推精度和性能，每一步都留下检查记录。这个过程熟练后，跑一个转换可能几分钟就完成，但省下来的是后续部署阶段几天甚至几周的排查时间。

如果只让我分享一条经验，那就是：onnx 导出成功不是终点，用 onnxruntime 跑通模型并对比精度，才算是真正完成了模型转换这一步。多花这五分钟做对比验证，比之后在目标设备上怀疑人生要划算得多。