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 版本关系密切。我用过不少版本组合，下面这张表是经过实际项目验证的相对稳定的搭配，供你参考：