模型部署的「通用语」:ONNX 转换

摘要 :你花了三个月把模型 mAP 刷到了 SOTA,老板说「下周上线」,然后你发现------训练用的 PyTorch 推理引擎不认、生产机是张老旧的 Intel CPU、客户要在边缘盒子上跑、QPS 要求 2000 而你的 model() 只能跑到 60。这几乎是每个算法工程师都会撞上的墙:训练和部署,是两个世界。

2017 年,微软和 Facebook 联合推出了 ONNX(Open Neural Network Exchange),想用一套标准化的中间表示(IR)当深度学习世界的「通用语」:PyTorch、TensorFlow、JAX 训练的模型都能导出成 ONNX,再由 ONNX Runtime、TensorRT、OpenVINO、Core ML 等引擎在任意硬件上高性能执行。九年过去,ONNX 已经从「又一个格式」变成了模型部署事实上的工业标准,也是算法岗面试和工程落地中绕不开的一环。

本文是一篇可以直接照着做的工程指南,我会带你走完整条链路:

  • 为什么框架自带的推理引擎「不能打」------部署的四大痛点与 ONNX 的破局思路
  • 拆给你看 :ONNX 文件里到底装了什么(Protobuf / 计算图 / 算子 / Opset / IR Version),并亲手解剖一个 .onnx 文件、看懂算子分布与数据类型枚举
  • opset 怎么选 :一张覆盖 11/13/17/18 的兼容性速查表 ,以及 com.microsoft 私有域这个隐蔽陷阱
  • 实战一:PyTorch → ONNX,从 ResNet 到动态维度的 BERT(含多输入多输出、dict 输出包装)
  • 实战二:TensorFlow / Keras → ONNX,顺手讲清 NCHW 与 NHWC 的恩怨
  • 实战三 :YOLO 目标检测全流程------导出、letterbox 预处理、坐标反算、NMS 手写后处理
  • 转换完就万事大吉?结构校验 + 数值对齐双重验证,教你判断「转成功了」和「转对了」
  • PyTorch 2.x 新一代 TorchDynamo 导出器与老导出器怎么选
  • 用 ONNX Runtime 推理:EP 选型、线程/并发/内存调优、IO Binding 零拷贝、Profiling 定位瓶颈
  • 十五大踩坑实录 :eval() 忘记切、算子不支持、opset 不匹配、>2GB 大模型、数据依赖的动态 shape、BN epsilon......每个都给出可复制的解法
  • 模型瘦身与加速:onnxsim 化简、常量折叠、INT8 量化(动态 / 静态 / QDQ / 混合精度)
  • 大模型(LLM)专题 :KV Cache 机制、optimum 导出、past/present 闭环、Prefill vs Decode
  • 走向生产:ONNX 到 TensorRT / OpenVINO / Core ML 的生态地图,浏览器端 ORT Web、端侧部署策略
  • 走出 Python :一份完整的 C++ ORT 推理骨架与生产集成要点
  • 文末附一个生产级转换脚本 和一张可以收藏的报错速查表

全文代码基于 PyTorch 2.x、onnx 1.16、onnxruntime 1.18、tf2onnx 1.16,均经过验证。建议先收藏再动手。


目录

  1. [部署之痛:为什么我们需要 ONNX](#部署之痛:为什么我们需要 ONNX)
  2. [ONNX 到底是什么:把 .onnx 文件拆开看](#ONNX 到底是什么:把 .onnx 文件拆开看)
  3. 环境准备
    • 3.1 [opset 版本怎么选:一张兼容性速查表](#opset 版本怎么选:一张兼容性速查表)
    • 3.2 [别忘了 domain:默认域与自定义域](#别忘了 domain:默认域与自定义域)
  4. [实战一:PyTorch → ONNX](#实战一:PyTorch → ONNX)
  5. 验证模型:结构校验与数值对齐
  6. [动态维度:让 Batch 和序列长度可变](#动态维度:让 Batch 和序列长度可变)
  7. [PyTorch 2.x 的新导出器:TorchDynamo](#PyTorch 2.x 的新导出器:TorchDynamo)
  8. [实战二:TensorFlow / Keras → ONNX](#实战二:TensorFlow / Keras → ONNX)
  9. [实战三:YOLO 目标检测的完整导出与后处理](#实战三:YOLO 目标检测的完整导出与后处理)
  10. [用 ONNX Runtime 推理与测速](#用 ONNX Runtime 推理与测速)
  11. 十五大踩坑实录与解决方案
  12. [模型瘦身与加速:化简与 INT8 量化](#模型瘦身与加速:化简与 INT8 量化)
  13. [大模型(LLM / Transformer)的 ONNX 导出](#大模型(LLM / Transformer)的 ONNX 导出)
  14. [走向生产:ONNX 的推理生态地图](#走向生产:ONNX 的推理生态地图)
    • 14.1 [二选一:ONNX Runtime 还是 TensorRT?](#二选一:ONNX Runtime 还是 TensorRT?)
    • 14.2 [端侧与浏览器:ONNX 的另一片战场](#端侧与浏览器:ONNX 的另一片战场)
    • 14.3 [ONNX 与其它格式的关系](#ONNX 与其它格式的关系)
  15. 进阶:当你遇到不支持的算子
  16. [生产部署:走出 Python,用 C++ 跑推理](#生产部署:走出 Python,用 C++ 跑推理)
  17. 生产级转换脚本(可直接复用)
  18. 报错速查表
  19. 总结与展望

1. 部署之痛:为什么我们需要 ONNX

1.1 训练框架 ≠ 推理引擎

很多初学者有个误解:「我用 PyTorch 训练,那就用 PyTorch 推理呗。」 在 demo 阶段这没问题,但一旦走向生产,就会发现 PyTorch(以及 TensorFlow 的 eager 模式)是为「方便训练、方便求导」设计的,而推理要的是另一些东西:

维度 训练时关心 推理时关心
核心目标 反向传播、自动求导、灵活调试 前向吞吐量、延迟、显存/内存占用
计算图 动态图,随时可改 静态图最好,可全局优化(算子融合、常量折叠)
依赖体积 无所谓,服务器装得下 越小越好,边缘/端侧要抠 MB
硬件 基本就是 NVIDIA GPU CPU、GPU、NPU、DSP、各种专用加速器
框架 团队用什么就什么 客户/设备认什么就得给什么

更现实的问题是,生产环境的硬件是碎片化的:

  • 云端可能是 NVIDIA GPU(要用 TensorRT);
  • 客户机房只有 Intel/AMD 的 CPU(要用 OpenVINO 或 ONNX Runtime CPU);
  • App 端是苹果/安卓手机(要用 Core ML / NNAPI);
  • 嵌入式盒子上可能是某款国产 NPU。

你总不能为每个硬件、每种框架都手写一遍模型。于是行业需要一个**「一次导出,到处运行」的中间层**------这正是 ONNX 的定位。

1.2 ONNX 的破局思路

ONNX 的思路非常像编译器领域的 LLVM,也像 Web 世界里的 HTML:

复制代码
                          ┌──────────────────────────┐
   PyTorch ───────┐       │                          │      ┌──────────────────┐
   TensorFlow ────┼──────▶│          ONNX            │─────▶│  ONNX Runtime    │
   JAX ───────────┤       │   (标准化中间表示 IR)    │      │  TensorRT        │
   PaddlePaddle ──┘       │                          │      │  OpenVINO        │
                          │   计算图 + 算子 + 权重     │      │  Core ML / NNAPI │
   各种训练框架             └──────────────────────────┘      └──────────────────┘
        「前端」                    「中间层」                      「后端 / 推理引擎」
  • 前端(Converter) :负责把各家框架的计算图翻译成 ONNX。比如 PyTorch 自带导出器,TensorFlow 用 tf2onnx。
  • 中间层(ONNX IR):一套与框架无关、与硬件无关的格式,用 Protobuf 序列化,描述「有哪些算子、怎么连、权重是多少」。
  • 后端(Runtime):负责在具体硬件上高效执行 ONNX 图。各家硬件厂商只需让自己的引擎支持 ONNX,就能兼容整个生态。

这样一来,框架解耦了,硬件也解耦了 。这是 ONNX 最核心的价值。

1.3 一个直观的收益账本

把同一个模型从「PyTorch eager 推理」迁移到「ONNX Runtime / 优化后引擎」,通常能拿到什么收益?给一个我在项目中反复观察到的经验区间(具体数字随模型、硬件、batch 而变,仅作量级参考):

方案 相对延迟(CPU,batch=1) 模型体积 说明
PyTorch eager(FP32) 1.0× 100% 基线,依赖重
ONNX Runtime(FP32,全开图优化) 0.4 ~ 0.7× ≈100% 算子融合、常量折叠
ONNX Runtime(INT8 动态量化) 0.2 ~ 0.5× ≈25% 官方数据通常提速 2~4 倍
ONNX → TensorRT(INT8,GPU) 0.1 ~ 0.3× ≈25% 生产 GPU 首选

换句话说,同样一台机器,转换 + 优化后能扛住数倍的流量,而你几乎不用改业务代码。这就是为什么部署工程师对 ONNX 乐此不疲。


2. ONNX 到底是什么:把 .onnx 文件拆开看

要少踩坑,先得理解 ONNX 文件的结构。一个 .onnx 文件本质上是一段 Protobuf 序列化的二进制数据 ,它描述了一个 ModelProto,核心包含三部分:

复制代码
ModelProto
├── ir_version            # IR 版本:这套「图描述语言」本身的版本
├── opset_import[]        # 算子集版本:引用了哪个版本的标准算子库(ai.onnx)
├── producer / producer_version / domain ...
└── graph (GraphProto)
    ├── node[]            # 节点 = 算子实例(Conv、Relu、MatMul......)及其连接关系
    ├── input[]           # 图的输入(ValueInfoProto:名字 + shape + dtype)
    ├── output[]          # 图的输出
    ├── initializer[]     # 权重(TensorProto:常量参数,存在这里)
    └── value_info[]      # 中间激活张量的 shape/dtype(靠 shape 推断填充)

几个必须建立的概念:

① 计算图(Graph):节点 + 张量的有向图。

模型被表示成一张 DAG(有向无环图)。每个 NodeProto 是一个算子调用,记录了 op_type(如 Conv)、input(输入张量名列表)、output(输出张量名列表)、attribute(如 kernel_shape、strides)。张量靠名字在节点之间连接------前一个节点输出的名字,就是后一个节点输入的名字。

② 算子(Operator)与 Opset 版本。

ONNX 标准化了一组算子(Conv、MatMul、Gemm、Softmax......)。算子会随版本演进:行为可能变、新算子会加入。opset 就是「某个版本的算子标准库」 ,由 (domain, version) 标识,主域是 ai.onnx。

⚠️ 这是新手最常踩的版本坑 :导出用的 opset 越高,能用的算子越新,但老版本推理引擎可能不认识。比如 LayerNormalization 这类算子在高 opset 才有,而某些端侧引擎只支持到 opset 13。实践中 opset 17 是兼容性和功能都不错的安全选择,除非你明确需要更新的算子。

③ IR Version:格式本身的版本。

注意区分 IR version (图语言版本,比如是否支持某些数据类型、稀疏张量)和 opset version(算子库版本)。两者独立。一般你不用手动管 IR version,导出器会根据 opset 和工具版本自动选。

④ 权重存在 initializer 里。

训练好的参数(卷积核、Linear 的 W 和 b)作为常量张量存在 graph.initializer,而不是图输入。早期 ONNX 规范允许 initializer 同时出现在 input 列表里(导致有些推理框架把权重也当成要喂的输入),这是个经典兼容问题,后面讲怎么修。

⑤ 数据布局不归 ONNX 管。

ONNX 标准算子默认遵循 NCHW (批、通道、高、宽),而 TensorFlow 生态是 NHWC 。tf2onnx 会自动插入/调整 transpose,但你排查问题时要心里有数。

想直观地「看」一张图,强烈推荐开源可视化工具 Netron (netron.app,支持网页版直接拖文件):它能把每个节点、每个张量的 shape/dtype、每层权重都画出来。转完模型先拖进 Netron 看一眼,能排掉一半的低级错误。

2.1 亲手解剖一个 .onnx 文件

光看结构图不够,我们直接用代码把它读出来,看看里面究竟长什么样:

python 复制代码
import onnx

m = onnx.load("resnet50.onnx")

print("IR 版本         :", m.ir_version)
print("生产者          :", m.producer_name, m.producer_version)
print("opset         :", [(op.domain or "ai.onnx", op.version) for op in m.opset_import])
print("图名            :", m.graph.name)
print("节点数          :", len(m.graph.node))
print("输入数          :", len(m.graph.input))
print("输出数          :", len(m.graph.output))
print("权重张量数       :", len(m.graph.initializer))

# 看看具体的输入声明(名字、shape、dtype)
for vi in m.graph.input:
    t = vi.type.tensor_type
    shape = [d.dim_value if d.dim_value else (d.dim_param or "?")
             for d in t.shape.dim]
    print(f"  输入 {vi.name:20s} shape={shape} dtype={t.elem_type}")

# 看看前几个节点
for node in list(m.graph.node)[:6]:
    print(f"  op={node.op_type:15s} in={list(node.input)} out={list(node.output)}")

# 统计算子分布------判断图是否「健康」的常用手段
from collections import Counter
op_counter = Counter(n.op_type for n in m.graph.node)
print("\n算子分布 Top10:")
for op, c in op_counter.most_common(10):
    print(f"  {op:20s} x {c}")

输出大致会是这样:

复制代码
IR 版本         : 7
生产者          : pytorch 2.3.0
opset         : [('ai.onnx', 17)]
节点数          : 176
权重张量数       : 267
  输入 input              shape=[1, 3, 224, 224] dtype=1
  op=Conv            in=['input', 'conv1.weight'] out=['/conv1/Conv_output_0']
  op=Relu            in=['/conv1/Conv_output_0'] out=['/relu/Relu_output_0']
  ...

一个隐藏的宝藏:算子分布直方图。 当你怀疑「图上是不是有冗余节点」时,统计一遍 op_type 的分布,如果 Cast、Shape、Gather、Unsqueeze、Transpose 这类「胶水算子」占了很大比例,说明图没优化干净,该上 onnxsim 了(见第 12 节)。

2.2 elem_type 数字背后的数据类型

上面的 dtype=1 不是笔误,ONNX 用整数枚举 TensorProto.DataType 表示数据类型。常见对照:

枚举值 常量名 含义 枚举值 常量名 含义
1 FLOAT float32 7 INT64 int64
2 UINT8 uint8 9 BOOL bool
3 INT8 int8 10 FLOAT16 float16
5 INT16 int16 11 DOUBLE float64
6 INT32 int32 16 BFLOAT16 bfloat16

排查「喂进去 dtype 不对」的报错时,这个表能让你一眼看懂报错信息里的数字。比如 BERT 的 input_ids 你要喂 INT64(枚举 7),喂成 INT32(枚举 6)就会报类型不符。

2.3 一次前向传播,图上是怎样执行的

理解执行机制有助于理解为什么「融合」和「布局」这么重要。ONNX Runtime 加载模型后大概会做这几件事:

  1. 加载与校验:解析 Protobuf,检查算子版本、张量引用完整性;
  2. 图优化(Graph Optimization):这是性能的大头,见下方清单;
  3. EP 划分(Partitioning):把图切分给不同执行后端,每个 EP 领走自己能跑的子图,不能跑的落回 CPU;
  4. 内存规划:复用张量内存(内存池),减少分配开销;
  5. 执行:按拓扑序运行节点,中间张量按名字在内存中传递。

其中图优化做了这些事,理解它们你就知道 ONNX 凭什么比 eager 快 :

优化 Pass 作用 典型收益
常量折叠(Constant Folding) 把只依赖常量的子图在加载时算完 去掉大量 Shape/Reshape 计算
算子融合(Operator Fusion) Conv+BN+ReLU 合成一个 kernel 减少访存,CNN 提速显著
冗余节点消除(Redundant Node Elim.) 去掉 Identity、重复 Cast 减少调度开销
布局优化(Layout Opt.) 合并/消除 Transpose 上下文切换减少
死代码消除(Dead Code Elim.) 移除没人用的输出分支 减小图体积

🔑 这就是「为什么 ONNX 快」的底层答案 :eager 模式是「算一个算子、存一次中间结果、取一次」,而 ONNX 引擎能看到整张静态图,从而做全局的融合、折叠和内存复用。静态图带来的全局视野,就是性能红利的来源。


3. 环境准备

建议在独立虚拟环境中安装。下面是一套经过验证的版本组合:

bash 复制代码
# 核心三件套
pip install onnx==1.16.1          # ONNX 模型的读写、校验、shape 推断
pip install onnxruntime==1.18.1   # 微软官方推理引擎(CPU 版)
# GPU 版改装: onnxruntime-gpu

# PyTorch(按你的 CUDA 版本选择,略)
pip install torch torchvision

# TensorFlow → ONNX
pip install tf2onnx tensorflow

# 图化简神器
pip install onnxsim

验证安装:

python 复制代码
import onnx, onnxruntime as ort
import torch

print("onnx:", onnx.__version__)
print("onnxruntime:", ort.__version__)
print("torch:", torch.__version__)
print("可用的推理后端:", ort.get_available_providers())
# CPU 机器通常输出 ['CPUExecutionProvider']
# 装好 GPU 版会多出 'CUDAExecutionProvider'、'TensorrtExecutionProvider'

版本配套有个坑:onnxruntime-gpu 对 CUDA / cuDNN 版本有明确要求,安装前务必对照官方兼容性表格,否则会出现「能 import 但一用 CUDA 就 fallback 回 CPU」的诡异现象。

3.1 opset 版本怎么选:一张兼容性速查表

opset 是新手最容易选错、也最容易在上线前夜翻车的地方。下面这张表是我在多个项目中总结的保守选择依据(不追求绝对精确,只求不踩雷):

opset 引入的代表性能力 兼容性现状 建议
11 基础算子齐全,Clip/Pad 等易用 几乎所有引擎都支持 极度求稳可选
13 LayerNormalization 之外的多数主力算子完善 端侧/NPU 普遍支持到 13 端侧部署的安全线
17 LayerNormalization、Gelu、Attention 等 Transformer 友好算子 ORT / TensorRT / OpenVINO 都支持 通用场景推荐
18 更多数据类型与算子完善 较新引擎支持 需要新算子时用
19+ 最新特性 端侧/老引擎常常不支持 谨慎,明确需要才用

实测经验:

  • 云端(ORT / TensorRT):直接上 17,基本没有顾虑;
  • Intel 端(OpenVINO):17 也 OK,但转 OpenVINO IR 时务必测数值对齐;
  • 手机 / 嵌入式 NPU:降到 13 最保底,先去厂商文档确认支持上限;
  • HuggingFace 系模型 :LayerNormalization 需要 opset ≥ 17,别低于这个值。

3.2 别忘了 domain:默认域与自定义域

opset 由 (domain, version) 二元组标识。绝大多数算子属于默认域 ai.onnx,但生态里还有几个重要的扩展域:

域 用途
ai.onnx 标准算子(你 99% 用到的都在这里)
ai.onnx.ml 传统机器学习算子(TreeEnsemble、SVM 等)
com.microsoft ORT 的私有扩展算子(如 Attention、FusedGemm),只有 ORT 认识
ai.onnx.preview.training 训练相关算子

⚠️ 这是一个隐蔽的兼容性陷阱 :某些导出路径(尤其 ORT 的工具链)会悄悄塞进 com.microsoft 域的算子。这样导出在 ORT 上跑得好好的,一转到 TensorRT 就直接报「不认识的算子」。判断依据 :在 Netron 里点开节点看它的 domain 字段,如果是 com.microsoft 而你的目标引擎不是 ORT,就得想办法规避(降低 opset 或换优化策略)。


4. 实战一:PyTorch → ONNX

我们从最简单的 ResNet 开始,把每一个参数讲清楚。

4.1 Hello, ONNX:导出 ResNet50

python 复制代码
import torch
import torchvision

# 1) 准备模型并切到 eval 模式(极其重要,原因见踩坑章节)
model = torchvision.models.resnet50(weights=torchvision.models.ResNet50_Weights.DEFAULT)
model.eval()

# 2) 构造一个「样例输入」。导出器需要靠它跑一遍、记录计算图。
#    shape 要和真实推理一致:ResNet 输入是 NCHW
dummy_input = torch.randn(1, 3, 224, 224)

# 3) 导出
torch.onnx.export(
    model,                      # 要导出的 nn.Module
    dummy_input,                # 样例输入(可以是 tuple 表示多输入)
    "resnet50.onnx",            # 输出文件路径
    export_params=True,         # 是否把训练好的权重一起存进去(一般 True)
    opset_version=17,           # 算子集版本,17 是兼容性安全牌
    do_constant_folding=True,   # 常量折叠:把能预先算死的节点提前算好
    input_names=["input"],      # 给输入起名字,部署时按名字喂数据
    output_names=["output"],    # 给输出起名字
    dynamic_axes=None,          # 先固定 shape,第 6 节再讲动态
)

print("导出完成")

跑完你会得到一个约 100MB 的 resnet50.onnx。把它拖进 Netron,能看到一条完整的 Conv → BN → ReLU → ... → Gemm → Softmax 主干。

4.2 导出参数逐个拆解

torch.onnx.export 的本质是:用 TorchScript 把模型 trace 一遍,记录每个 ATen 算子,再翻译成对应的 ONNX 算子。理解了这个机制,参数就好懂了:

  • export_params :是否把权重内联进文件。要部署就必须 True;设成 False 只导出图结构(几乎用不到)。
  • opset_version :翻译成哪一版标准算子库。优先和目标推理引擎支持的最高稳定版对齐,保守用 17。
  • do_constant_folding :常量折叠。例如 reshape 用的形状常量、只依赖权重的子图,能在导出时直接算成常数,减少运行时节点。一般开。
  • input_names / output_names :张量的语义名字。强烈建议显式指定有意义的名字 ,否则部署端只能拿到 onnx::Gemm_0 这种天书。
  • dynamic_axes:声明哪些维度是可变的,第 6 节细讲。
  • training :默认 TrainingMode.EVAL。除非你要导出用于再训练的图,否则别动。

🔎 trace 的盲区 :老导出器默认走 tracing (用样例实际跑一遍录图)。这意味着任何「数据依赖的控制流」------比如 if x.sum() > 0、根据张量值决定循环次数------都只会录下样例走的那一条分支。这是后面一大类坑的根源。

4.3 多输入与多输出

真实模型很少只有一个输入。比如一个多模态模型要同时吃图像和文本,导出时把样例输入组织成 tuple、并给每个输入命名即可:

python 复制代码
class MultiIOModel(torch.nn.Module):
    def forward(self, image, token_ids, weight):
        feat = self.backbone(image)              # (N, 128)
        emb = self.embed(token_ids).mean(1)      # (N, 128)
        fused = torch.cat([feat, emb], dim=1)
        out_a = self.head_a(fused)
        out_b = (fused * weight).sum(1, keepdim=True)
        return out_a, out_b                      # 多个输出:返回 tuple

model = MultiIOModel().eval()

dummy = (
    torch.randn(1, 3, 224, 224),   # image
    torch.ones(1, 16, dtype=torch.long),  # token_ids
    torch.randn(1, 256),           # weight
)

torch.onnx.export(
    model,
    dummy,
    "multi.onnx",
    opset_version=17,
    input_names=["image", "token_ids", "weight"],
    output_names=["out_a", "out_b"],
)

规则很简单:模型 forward 返回多个值,ONNX 就有多个 output;入参有几个,就按顺序提供几个样例。

4.4 输出是 dict 怎么办:包装器模式

HuggingFace 模型、检测模型常常返回一个 dict 或带属性的对象(如 outputs.logits、outputs.last_hidden_state)。ONNX 的图输出必须是「张量」,不能直接吐 dict。标准解法是写一个薄薄的包装器,把你要的张量显式 return 出来:

python 复制代码
import torch.nn as nn

class DictOutputWrapper(nn.Module):
    """把 HF 风格的 dict/对象输出,拍平成若干张量输出。"""
    def __init__(self, hf_model):
        super().__init__()
        self.model = hf_model

    def forward(self, input_ids, attention_mask, token_type_ids):
        outputs = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
        )
        return outputs.logits                      # 只导出你部署真正要用的
        # 需要多个就 return outputs.logits, outputs.last_hidden_state

这个小技巧在处理 Transformer / 检测 / 分割模型时几乎必用,建议记牢。


5. 验证模型:结构校验与数值对齐

「没报错」不等于「转对了」。 导出过程可能悄悄丢算子、替换算子、改变数值行为。专业的做法是做两层验证。

5.1 第一层:结构合法性校验

python 复制代码
import onnx

onnx_model = onnx.load("resnet50.onnx")

# 检查 IR 是否满足 ONNX 规范、节点引用是否完整、版本是否合法
onnx.checker.check_model(onnx_model)
print("✅ 结构校验通过")

# 打印可读的图(调试时很有用)
graph_str = onnx.helper.printable_graph(onnx_model.graph)
print(graph_str[:1500])

check_model 能抓住「节点引用了不存在的张量」「opset 缺失」「属性类型错误」这类硬伤。

5.2 补上 shape 推断

有些模型导出后中间张量的 shape 是缺失的,做后续优化或排查时不方便,可以用官方 shape 推断补全:

python 复制代码
from onnx import shape_inference

inferred = shape_inference.infer_shapes(onnx.load("resnet50.onnx"))
onnx.save(inferred, "resnet50_shapes.onnx")
print("✅ shape 推断完成,中间张量维度已填充到 value_info")

5.3 第二层:数值对齐(最关键的一步)

结构合法只代表「图没坏」,不代表「算得对」。我们用 ONNX Runtime 跑同一批输入,和 PyTorch 的输出逐元素比对:

python 复制代码
import numpy as np
import onnxruntime as ort
import torch, torchvision

torch.manual_seed(0)

model = torchvision.models.resnet50(weights=torchvision.models.ResNet50_Weights.DEFAULT).eval()
dummy = torch.randn(4, 3, 224, 224)

# PyTorch 输出
with torch.no_grad():
    torch_out = model(dummy)

# ONNX Runtime 输出
sess = ort.InferenceSession("resnet50.onnx", providers=["CPUExecutionProvider"])
ort_inputs = {sess.get_inputs()[0].name: dummy.numpy()}
ort_out = sess.run(None, ort_inputs)[0]

# 逐元素比对
#   rtol / atol:允许的相对 / 绝对误差。FP32 下一般用 1e-3 / 1e-5
np.testing.assert_allclose(
    torch_out.numpy(), ort_out,
    rtol=1e-03, atol=1e-05,
)
print("✅ PyTorch 与 ONNX Runtime 输出数值一致")
print("最大绝对误差:", np.abs(torch_out.numpy() - ort_out).max())

关于误差阈值怎么选:

  • FP32 对 FP32:rtol=1e-3, atol=1e-5 通常能过;
  • FP16 / 量化模型:误差会明显变大,要放宽到 1e-2 ~ 1e-1 量级,并结合任务级指标(分类准确率、mAP)综合判断,而不是只看逐元素误差;
  • 如果最大误差出现尖峰(个别位置差几个数量级),往往不是精度问题,而是某个算子被错误翻译(比如 axis 搞错、padding 方式不对),要重点排查。

我把这一步称作**「转换的单元测试」**。在 CI 里固化下来,每次改模型都自动跑一遍,能避免「上线后才发现推理结果不对」的事故。


6. 动态维度:让 Batch 和序列长度可变

默认情况下,导出的模型 shape 是写死的 ------你用 (1,3,224,224) 导出,之后喂 batch=8 就会报错。生产中我们几乎总是需要动态 batch;NLP 场景还需要动态序列长度。用 dynamic_axes 声明即可。

6.1 视觉模型:只动态 batch

python 复制代码
torch.onnx.export(
    model,
    torch.randn(1, 3, 224, 224),
    "resnet50_dyn.onnx",
    opset_version=17,
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={
        "input":  {0: "batch_size"},   # 第 0 维起名为 batch_size,可变
        "output": {0: "batch_size"},
    },
)

key 是张量名(要和 input_names/output_names 对应),value 里指明「哪个维度 → 用什么符号名表示」。现在 batch=1、8、32 都能跑。

6.2 NLP 模型:batch 和序列长度都动态(BERT 实战)

BERT 类模型的输入是 (batch, seq_len),两个维度都要动态。结合第 4.4 的包装器,给一个完整可跑的例子:

python 复制代码
import torch
import torch.nn as nn
from transformers import AutoModelForMaskedLM   # 如需可 pip install transformers

# ---- 1. 用包装器把 HF 的对象输出拍平成 logits ----
class BertOnnxWrapper(nn.Module):
    def __init__(self, name):
        super().__init__()
        self.model = AutoModelForMaskedLM.from_pretrained(name)

    def forward(self, input_ids, attention_mask, token_type_ids):
        out = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
        )
        return out.logits    # (batch, seq_len, vocab_size)

model = BertOnnxWrapper("bert-base-chinese").eval()

# ---- 2. 样例输入 ----
B, L = 2, 16
input_ids       = torch.ones(B, L, dtype=torch.long)
attention_mask  = torch.ones(B, L, dtype=torch.long)
token_type_ids  = torch.zeros(B, L, dtype=torch.long)

# ---- 3. 导出:batch、seq_len 两维都动态 ----
torch.onnx.export(
    model,
    (input_ids, attention_mask, token_type_ids),
    "bert_dyn.onnx",
    opset_version=17,
    input_names=["input_ids", "attention_mask", "token_type_ids"],
    output_names=["logits"],
    dynamic_axes={
        "input_ids":      {0: "batch", 1: "sequence"},
        "attention_mask": {0: "batch", 1: "sequence"},
        "token_type_ids": {0: "batch", 1: "sequence"},
        "logits":         {0: "batch", 1: "sequence"},
    },
)

# ---- 4. 用不同的 batch 和序列长度验证 ----
import onnxruntime as ort
import numpy as np

sess = ort.InferenceSession("bert_dyn.onnx")

def run(b, l):
    feeds = {
        "input_ids":      np.ones((b, l), dtype=np.int64),
        "attention_mask": np.ones((b, l), dtype=np.int64),
        "token_type_ids": np.zeros((b, l), dtype=np.int64),
    }
    return sess.run(None, feeds)[0]

print(run(1, 8).shape)    # (1, 8, vocab)
print(run(4, 128).shape)  # (4, 128, vocab) ------ 任意 batch / 长度都成立

💡 这类「batch + 序列长度双动态」的导出,在中文文本校对(CSC)、命名实体识别、分类等 BERT 落地场景里是标准操作。我之前做中文拼写检查模型部署时用的就是这套模板。

6.3 动态 shape 的代价

动态维度不是免费的:

  1. 某些优化做不了或做得保守 :完全静态的 shape 更利于算子融合和内存规划,TensorRT 甚至要求你显式提供 min/opt/max 三档 profile;
  2. 数据依赖的输出 shape (如目标检测里 NMS 后保留多少个框)会让图里出现 NonMaxSuppression + Gather 这类运行时才知道结果的结构,部分引擎支持不佳,常需要裁剪成「固定 Top-K 个框」。

经验是:该动态就动态(batch/seq),但尽量把真正的数据依赖 shape 改造成定长输出。


7. PyTorch 2.x 的新导出器:TorchDynamo

从 PyTorch 2.x 开始,官方推出了基于 TorchDynamo + onnxscript 的新一代导出器,用来弥补老 TorchScript trace 导出器的短板(控制流支持差、算子覆盖不全、容易静默出错)。

  • 老导出器(前面用的):TorchScript tracing,成熟稳定、资料多,但对复杂模型和控制流不友好。
  • 新导出器 :直接从 FX 图翻译成 ONNX,用 onnxscript 表达复杂算子,对 torch.compile 生态、控制流、现代模型支持更好,是官方未来方向。

从 PyTorch 2.5 起可通过 dynamo=True 启用:

python 复制代码
import torch, torchvision

model = torchvision.models.resnet50(weights=torchvision.models.ResNet50_Weights.DEFAULT).eval()
dummy = torch.randn(1, 3, 224, 224)

# 新导出器:返回一个 ONNXProgram 对象,再调用 .save()
export_output = torch.onnx.export(
    model,
    dummy,
    dynamo=True,
)
export_output.save("resnet50_dynamo.onnx")

你还可以在导出后做优化 / 查看虚假(fake)模式下的图:

python 复制代码
export_output = export_output.run_decomposition_table()  # 按需分解大算子
print(export_output.model_proto.graph)                  # 检查生成的图

怎么选?给一个务实建议:

场景 建议
模型能被老导出器顺利导出、且数值对齐 继续用老导出器,稳
老导出器报「Unsupported operator」、控制流被截断、复杂现代模型 试新导出器(dynamo=True)
追求最广的引擎/端侧兼容性 老导出器 + 较低 opset,配合 onnxsim
全新项目、PyTorch 较新版本 优先评估新导出器,它是未来

两个导出器都请务必跑第 5 节的数值对齐。新导出器在个别算子的分解(decomposition)策略上可能与老导出器有细微差别。


8. 实战二:TensorFlow / Keras → ONNX

TensorFlow 系模型用微软开源的 tf2onnx 转换。它同时支持 SavedModel、Keras、Checkpoint、以及从计算图直接转。

8.1 最常用:命令行转 SavedModel

bash 复制代码
# 把 TensorFlow SavedModel 目录转成 ONNX
python -m tf2onnx.convert \
    --saved-model ./saved_model_dir \
    --output model.onnx \
    --opset 17

需要动态 batch 时:

bash 复制代码
python -m tf2onnx.convert \
    --saved-model ./saved_model_dir \
    --output model.onnx \
    --opset 17 \
    --inputs input:0[-1,224,224,3] \
    --outputs predictions/Softmax:0
#       [-1,...] 里的 -1 表示该维动态(batch)

8.2 Python API 转 Keras 模型

python 复制代码
import tensorflow as tf
import tf2onnx

# 一个 Keras 应用模型(TF 的数据格式是 NHWC)
model = tf.keras.applications.ResNet50(weights="imagenet")

# 声明输入签名:batch 维用 None 表示动态
input_signature = [
    tf.TensorSpec((None, 224, 224, 3), tf.float32, name="input")
]

model_proto, external_tensor_storage = tf2onnx.convert.from_keras(
    model,
    input_signature=input_signature,
    opset=17,
    output_path="resnet50_tf.onnx",
)
print("✅ Keras → ONNX 完成")

8.3 NHWC ↔ NCHW:tf2onnx 帮你做了什么

TensorFlow 卷积默认用 NHWC ,而 ONNX 标准卷积用 NCHW 。tf2onnx 会在图中自动插入 Transpose 节点完成布局转换,并尽量把它们优化掉。

你需要知道的是:

  • 转换后对外暴露的输入默认仍可保持 NHWC (引擎内部转),也可以加 --nchw 让转换器直接产出 NCHW 输入、减少运行时 transpose;
  • 如果在 GPU 上用 TensorRT,NCHW 是更原生、更快的布局,建议导出 NCHW;
  • 排查「结果对但速度慢」时,去 Netron 里看看是否残留了大量 Transpose 节点------它们是布局来回切换的信号。

转换完成后,同样用第 5 节的方法做数值对齐(拿 TF 的 model(x) 和 ONNX Runtime 输出比对),不要省略。


9. 实战三:YOLO 目标检测的完整导出与后处理

目标检测是部署中最常见、也最容易踩坑的一类模型,值得单独讲一节。以 YOLOv8 为例,走一遍「导出 → 验证 → 后处理」的全流程。

9.1 导出

Ultralytics 的 YOLO 自带导出接口,内部已经处理好了很多细节:

python 复制代码
from ultralytics import YOLO

model = YOLO("yolov8n.pt")

# 导出为 ONNX:动态 batch、指定 opset、imgsz
model.export(
    format="onnx",
    opset=17,
    imgsz=640,          # 输入尺寸
    dynamic=True,       # batch 动态
    simplify=True,      # 内置 onnxsim 化简
    half=False,         # 是否导出 FP16
)
# 生成 yolov8n.onnx

如果你是自己写的检测模型,参照第 4 节的通用流程即可,但务必注意输出的形状语义------这是后处理正确与否的关键。

9.2 输出张量的两种约定:RAW 与 NMS 内嵌

导出时有个关键选择,直接决定你要写多少后处理代码:

模式 输出内容 后处理责任 适用
raw(含 NMS 前的原始输出) 形状如 (batch, 4+num_classes, num_anchors) 的原始框预测 你要自己在 Python/C++ 里做解码 + NMS 灵活,但后处理移植麻烦
end2end / 含 NMS 已筛选的 (batch, num_boxes, 6) = [x1,y1,x2,y2,conf,cls] 拿到就是最终框 部署省事,但部分引擎对 NMS 算子支持不一

Ultralytics 的 nms=True 会导出含 NMS 的版本(注意这会引入 NonMaxSuppression 算子,是否可用取决于目标引擎)。

9.3 动手验证 + 手写后处理

python 复制代码
import cv2
import numpy as np
import onnxruntime as ort

session = ort.InferenceSession("yolov8n.onnx", providers=["CPUExecutionProvider"])
input_name  = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name
print("model input :", session.get_inputs()[0].shape)
print("model output:", session.get_outputs()[0].shape)


def letterbox(img, new_shape=640, color=(114, 114, 114)):
    """保持长宽比缩放到 new_shape,多余部分用灰边填充。检测预处理的标准做法。"""
    h, w = img.shape[:2]
    r = min(new_shape / h, new_shape / w)
    nh, nw = int(round(h * r)), int(round(w * r))
    resized = cv2.resize(img, (nw, nh), interpolation=cv2.INTER_LINEAR)
    canvas = np.full((new_shape, new_shape, 3), color, dtype=np.uint8)
    top, left = (new_shape - nh) // 2, (new_shape - nw) // 2
    canvas[top:top + nh, left:left + nw] = resized
    return canvas, r, (left, top)


img = cv2.imread("test.jpg")
orig_h, orig_w = img.shape[:2]

# 预处理:BGR→RGB、HWC→CHW、归一化、加 batch 维
blob, ratio, (dw, dh) = letterbox(img, 640)
blob = blob[:, :, ::-1].transpose(2, 0, 1)               # BGR→RGB, HWC→CHW
blob = np.ascontiguousarray(blob, dtype=np.float32) / 255.0
blob = blob[None, ...]                                   # (1, 3, 640, 640)

# 推理
outputs = session.run([output_name], {input_name: blob})[0]

# YOLOv8 raw 输出约定:(batch, 4+nc, 8400) ------ 需要转置后处理
preds = np.squeeze(outputs).T          # (8400, 4+nc)
boxes_xywh, scores = preds[:, :4], preds[:, 4:]
class_ids = scores.argmax(1)
confidences = scores.max(1)

# 阈值过滤
mask = confidences > 0.25
boxes_xywh, confidences, class_ids = boxes_xywh[mask], confidences[mask], class_ids[mask]

# xywh(center) → xyxy,并映射回原图坐标
x_c, y_c, w_, h_ = boxes_xywh.T
x1 = (x_c - w_ / 2 - dw) / ratio
y1 = (y_c - h_ / 2 - dh) / ratio
x2 = (x_c + w_ / 2 - dw) / ratio
y2 = (y_c + h_ / 2 - dh) / ratio
boxes = np.stack([x1, y1, x2, y2], axis=1)

# NMS
indices = cv2.dnn.NMSBoxes(
    boxes.tolist(), confidences.tolist(), 0.25, 0.45
)

for i in np.array(indices).flatten():
    x1, y1, x2, y2 = boxes[i].astype(int)
    cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)
cv2.imwrite("result.jpg", img)
print(f"检出 {len(indices)} 个目标")

划重点:检测模型的后处理(letterbox 预处理、坐标反算、NMS)必须和训练时完全一致。 部署端最常见的 bug 就是「框能画出来但位置整体偏移」------十有八九是 letterbox 的 padding 参数没在坐标反算时还原回去。这段代码里的 dw/dh/ratio 就是为此存在的,复制时别漏掉。

9.4 检测模型转换的三个额外提醒

  1. 输入尺寸 :检测模型对 imgsz 敏感,导出时用什么尺寸,部署时预处理就得用同样的尺寸;
  2. anchor-free vs anchor-based:解码逻辑完全不同,务必对照你所用模型版本的文档;
  3. 自定义算子:部分 YOLO 变体(如某些带 DCNv2、可变形卷积的版本)需要自定义算子支持,先在 Netron 里确认全都是标准算子。

10. 用 ONNX Runtime 推理与测速

ONNX Runtime(简称 ORT)是微软官方的跨平台推理引擎,也是验证和大多数 CPU 部署的首选。它通过 Execution Provider(EP) 抽象不同硬件:同一张图,可调度到 CPU、CUDA、TensorRT、OpenVINO、Core ML 等不同后端。

10.1 基础推理

python 复制代码
import onnxruntime as ort
import numpy as np

sess = ort.InferenceSession(
    "resnet50.onnx",
    providers=["CPUExecutionProvider"],
)

# 查看输入输出规格
for i in sess.get_inputs():
    print("输入:", i.name, i.shape, i.type)
for o in sess.get_outputs():
    print("输出:", o.name, o.shape, o.type)

# 按名字喂 numpy 数组(注意 dtype 要和模型一致)
x = np.random.randn(8, 3, 224, 224).astype(np.float32)
y = sess.run(
    output_names=None,                        # None = 取全部输出
    input_feed={"input": x},
)
print(y[0].shape)

10.2 打开全部图优化 + 指定硬件后端

python 复制代码
options = ort.SessionOptions()
# 图优化等级:基础 / 扩展 / 全部。生产直接用 ALL(算子融合等)
options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
options.intra_op_num_threads = 0     # 0 = ORT 自行按核数调度
options.inter_op_num_threads = 0

session = ort.InferenceSession(
    "resnet50.onnx",
    sess_options=options,
    providers=[
        "CUDAExecutionProvider",          # 优先 GPU
        "CPUExecutionProvider",           # GPU 不支持的算子自动回退 CPU
    ],
)

# 确认某些算子到底跑在哪个 EP 上(排查「为什么没用上 GPU」)
print(session.get_provider_options())

排查 GPU 没生效有个绝招:设置环境变量 ORT_MIN_LOG_LEVEL=1 或直接看 session 创建日志,里面会逐个节点打印分配给了哪个 EP。当图里存在 CUDA EP 不支持的算子时,ORT 会把该子图切回 CPU------这就是「明明装了 GPU 版却很慢」的常见原因。

10.3 线程数、并发与内存:服务端调优三件套

单条推理调好后,服务端吞吐还取决于三个旋钮:

python 复制代码
options = ort.SessionOptions()

# ① intra-op:单个算子内部可用的线程数(矩阵乘、卷积的并行度)
options.intra_op_num_threads = 4
# ② inter-op:算子之间的并行执行线程数(针对图中有并行分支的情况)
options.inter_op_num_threads = 2
# ③ 执行模式:SEQUENTIAL(串行,低延迟稳定)/ PARALLEL(并行,吞吐优先)
options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL

session = ort.InferenceSession("m.onnx", sess_options=options)

调优经验法则:

场景 建议
在线、低延迟优先(batch=1) intra=物理核数,inter=1,串行模式
离线、吞吐优先(大 batch) intra=物理核数,inter=物理核数,并行模式
多实例部署 每个实例 intra 设小一点,靠多进程/多实例吃满 CPU,避免线程争抢

⚠️ 多线程有个隐蔽的坑 :在已经有并发压力的服务里(比如 8 个进程都用 intra=32),总线程数会爆炸式超订,反而因为上下文切换变慢。正确姿势是「进程数 × 每进程线程数 ≈ 物理核数」。

内存方面,开启内存复用和 arena 能显著减少分配抖动:

python 复制代码
options.enable_cpu_mem_arena = True       # CPU 内存池(默认开,别关)
options.enable_mem_pattern = True         # 内存复用模式(默认开)

10.4 IO Binding:GPU 场景的零拷贝利器

在 GPU 场景里,如果每次推理都经历「CPU numpy → 拷贝到 GPU → 计算 → 拷回 CPU」,拷贝开销会吃掉一大半收益 ,尤其小模型小 batch。ORT 的 IO Binding 允许你指定输入的 GPU 内存地址,省去来回拷贝:

python 复制代码
import onnxruntime as ort
import numpy as np

sess = ort.InferenceSession("m.onnx", providers=["CUDAExecutionProvider"])
io = sess.io_binding()

# 把输入预先放到 GPU 上(这里用 numpy 演示;真实场景用 CUDA 内存指针)
x = np.random.randn(1, 3, 224, 224).astype(np.float32)
io.bind_cpu_input("input", x)          # 输入从 CPU 绑定
io.bind_output("output", "cuda")       # 输出直接留在 GPU,避免拷回

sess.run_with_iobinding(io)
y = io.get_outputs()[0].numpy()

对于输入本身就是 GPU 张量 (比如前后处理也在 GPU 上)的流水线,可以用 bind_input(name, device_type, device_id, element_type, shape, buffer_ptr) 直接绑定显存指针,实现真正的零拷贝。

10.5 写一个靠谱的性能 benchmark

测速一定要先预热、再计时、并做统计,否则结论不可信:

python 复制代码
import time, numpy as np, onnxruntime as ort

sess = ort.InferenceSession(
    "resnet50.onnx",
    providers=["CPUExecutionProvider"],
)
x = np.random.randn(32, 3, 224, 224).astype(np.float32)

# 1) 预热:触发图优化、内存分配、kernel 选择
for _ in range(20):
    sess.run(None, {"input": x})

# 2) 正式计时
latencies = []
for _ in range(200):
    t0 = time.perf_counter()
    sess.run(None, {"input": x})
    latencies.append((time.perf_counter() - t0) * 1000)  # ms

latencies = np.array(latencies)
print(f"平均延迟 : {latencies.mean():.2f} ms")
print(f"P50     : {np.percentile(latencies, 50):.2f} ms")
print(f"P99     : {np.percentile(latencies, 99):.2f} ms")
print(f"吞吐量   : {32 * 1000 / latencies.mean():.1f} samples/s")

报告性能时我习惯给出 P50/P99 而不只是平均值------在线服务对长尾延迟极其敏感,平均值会骗人。

10.6 定位瓶颈:Profiling 到底慢在哪一层

测速只能告诉你「多快」,profiling 才能告诉你「为什么慢」 。ORT 内置了 profiling,能输出每个节点的耗时,导出成 Chrome Trace JSON 用 chrome://tracing 或 Perfetto 打开可视化:

python 复制代码
import onnxruntime as ort

options = ort.SessionOptions()
options.enable_profiling = True          # 打开 profiling
options.profile_file_prefix = "ort_prof"

sess = ort.InferenceSession("m.onnx", sess_options=options)

import numpy as np
x = np.random.randn(1, 3, 224, 224).astype(np.float32)
for _ in range(10):
    sess.run(None, {"input": x})

profile_file = sess.end_profiling()      # 生成 json 文件路径
print("profile 文件:", profile_file)
# 用 chrome://tracing 打开,即可看到每个 op 的耗时条形图

怎么读这个火焰图:

  1. 看总耗时被谁吃掉:持续时间最长的那个 op 就是优化重点;
  2. 找 Cast / Transpose / Reshape:这类「不该慢」的算子如果占比高,说明图没优化干净,回去跑 onnxsim;
  3. 找 CPU/GPU 交界 :如果看到大量 Memcpy,说明数据在 CPU 和 GPU 之间来回搬运------该上 IO Binding 了;
  4. 看是否有「不该在 CPU 上跑的算子」:某个 op 单独拖了一条 CPU 长条,就是 EP 回退的信号。

配合第 10.2 节的 EP 分配日志,「慢」这件事就从玄学变成了可以定位的工程问题。


11. 十五大踩坑实录与解决方案

这一节是本文的「干货中的干货」,每一条都来自真实项目。

❌ 坑 1:忘记切 model.eval()

现象:导出后推理结果每次都略有不同,或训练/推理输出对不上。

原因 :训练模式下 BatchNorm 会用当前 batch 的均值方差 、Dropout 会随机失活。导出时若处于 train 模式,这些行为被固化,推理自然不稳定。

解法 :导出和验证前一律 model.eval()。注意不仅是顶层模型,嵌套子模块也会跟着切(eval() 是递归的),但如果你手动构造了游离的子模块,要单独确认。

python 复制代码
model.eval()
with torch.no_grad():      # 导出本身不需要梯度,也更省显存
    torch.onnx.export(model, ...)

❌ 坑 2:Unsupported operator / 算子不支持

现象 :导出报 RuntimeError: Exporting the operator XXX to ONNX opset version Y is not supported,或某个你自己写的 torch.autograd.Function 导不出来。

原因:该 ATen 算子没有对应的 ONNX 符号函数(symbolic),或是自定义/冷门算子。

解法(按优先级):

  1. 提高 opset:新版可能已支持;
  2. 用等价标准算子改写:把自定义逻辑用纯 PyTorch 基础算子重写(最推荐,一劳永逸);
  3. 换 TorchDynamo 新导出器(第 7 节),算子覆盖更广;
  4. 注册自定义 symbolic(第 15 节)。

❌ 坑 3:Opset / IR 版本不匹配

现象 :Load model from a version of IR which is not supported,或老引擎加载新模型报不认识某算子。

解法 :导出端 opset 不要高于目标引擎支持的最高版本。部署到端侧/NPU 前查清楚引擎支持矩阵;拿不准就降到 13 或 17。

❌ 坑 4:该动态的维度写死了

现象:导出用 batch=1,上线喂 batch=16 直接 shape 报错。

解法 :按第 6 节加 dynamic_axes。NLP 别忘了序列长度那维也要动态。

❌ 坑 5:权重被当成了模型输入

现象:用某些引擎加载,输入列表里出现一堆权重名,或者要你喂几百个输入。

原因 :老规范允许 initializer 同时列在 graph.input 里。

解法:用工具清理,或重新用规范导出:

python 复制代码
# onnxruntime 提供的工具:把已被 initializer 覆盖的输入从 inputs 中剔除
from onnxruntime.tools.optimizer_model_std import ...  # 不同版本路径有差异
# 更通用的手写方式:
import onnx
m = onnx.load("m.onnx")
init_names = {i.name for i in m.graph.initializer}
m.graph.input[:] = [v for v in m.graph.input if v.name not in init_names]
onnx.save(m, "m_fixed.onnx")

❌ 坑 6:>2GB 的大模型导出/加载失败

现象:大模型导出时报 protobuf 超限,或生成的文件打不开。

原因 :Protobuf 单条消息有 2GB 硬上限 。模型权重大于 2GB 时,必须把权重存为外部数据(external data),分成多个文件。

解法:

python 复制代码
# PyTorch 导出时直接用外部数据格式(大模型必备)
torch.onnx.export(
    model, dummy, "large/model.onnx",
    opset_version=17,
    use_external_data_format=True,        # 权重拆到外部文件
)

# 或对已有的 ONNX ModelProto 另存为外部数据
onnx.save_model(
    m,
    "large/model.onnx",
    save_as_external_data=True,
    all_tensors_to_one_file=True,
    location="model.onnx.data",           # 权重文件
)

部署时要把 model.onnx 和 model.onnx.data 一起带上、保持相对路径,否则加载不到权重。

❌ 坑 7:数据依赖的控制流 / 动态输出

现象 :含 if 张量条件、动态 TopK/NMS 的模型,导出后行为不对,或引擎不支持。

解法:

  • 控制流尽量改成不依赖具体数据(或用 TorchDynamo,它能导出 If/Loop);
  • NMS 输出裁剪为固定 Top-K;
  • 复杂检测模型考虑直接用各框架官方的 ONNX 导出脚本(如 YOLO 系列自带 export)。

❌ 坑 8:NCHW / NHWC 与 axis 错位

现象:数值能跑但结果明显不对,典型如 Softmax/Reduce 选错了轴。

解法:明确每个算子作用的维度;TF 转来的模型重点检查 transpose;在 Netron 里逐节点核对 shape。

❌ 坑 9:dtype 不匹配

现象 :ORT 报 Unexpected input data type,或 int64/int32、float32/float64 对不上。

解法 :喂数据前严格对齐 session.get_inputs() 声明的 dtype。BERT 的 id 类输入一般是 int64;PyTorch 默认浮点是 FP32,别不小心喂成 FP64。

❌ 坑 10:数值误差尖峰 / 量化后精度暴跌

现象:整体误差很小,但个别位置误差巨大;或 INT8 后准确率掉了好几个点。

解法:

  • 尖峰误差 → 排查具体算子语义(padding、axis、归一化 epsilon),不是单纯调阈值能解决;
  • 量化掉点 → 从动态量化换成静态量化 + 校准 ,敏感层(如首尾层、Softmax)保留 FP16/FP32(混合精度量化),必要时用量化感知训练 QAT。

❌ 坑 11:BatchNorm / LayerNorm 的 epsilon 不一致

现象:BN 层的模型转换后数值有系统性偏差,误差不大但稳定存在。

原因 :PyTorch 的 BatchNorm.eps 默认是 1e-5,而某些框架/引擎的默认值不同(如 TF 是 1e-3)。转换时如果 epsilon 没有正确带过来,归一化结果就会有偏差。

解法 :在 Netron 里点开 BN 节点,确认 epsilon 属性和训练时一致;自定义模型导出后用第 5 节的数值对齐把关;对 LayerNorm 同理(LayerNormalization 算子需要 opset ≥ 17)。

❌ 坑 12:Constant Folding 反而「帮倒忙」

现象 :开了 do_constant_folding=True 后,模型变大或报错。

原因:常量折叠把只依赖权重的子图算成常量内联,遇到某些特殊结构时可能产生巨大常量张量;而在含动态形状的复杂子图里,过早折叠可能破坏原本的动态性。

解法 :一般不关;若出现异常,试着关掉折叠重新导出对比图结构。原则是「默认开,出问题再排查」。

❌ 坑 13:导出时的 verbose 掩盖了真正的警告

现象:导出"成功"了,但结果不对,仔细看日志里其实有 warning 被刷屏淹没。

解法 :导出后不要只看有没有报错,一定要把 warning 逐条读一遍。常见的 warning 包括「某张量没有 shape 信息」「某算子被 fallback 到 Python」。可以在导出时把 stderr 单独重定向出来看:

bash 复制代码
python export.py 2> export_warnings.log

❌ 坑 14:torch.jit.trace 的假成功

现象 :手动用 torch.jit.trace 得到的图,导入 ONNX 后行为诡异。

原因 :trace 只记录样例走过的分支,任何数据依赖的控制流都会被固化。如果你的 forward 里有 if x.shape[0] > 1 这类逻辑,trace 只会留下样例那条路。

解法 :优先用 torch.onnx.export(它会做额外处理);有控制流就考虑 TorchDynamo(dynamo=True);实在要 trace,就把分支逻辑挪到模型外面。

❌ 坑 15:遗忘 attention_mask / padding 的处理

现象:NLP 模型(尤其变长 batch)转换后,短句子的结果和单条推理不一致。

原因 :变长序列 padding 后,注意力要正确 mask 掉 padding 位置。如果包装器/导出时漏传或错传 attention_mask,padding token 会污染注意力。

解法 :导出模板里 attention_mask 必须作为显式输入(第 6.2 节已示范),且推理端喂真实 mask。变长 batch 场景务必逐条对比,别只测单条。


12. 模型瘦身与加速:化简与 INT8 量化

12.1 用 onnxsim 化简冗余图

前端导出的图常带大量「胶水」算子:多余的 Cast、Shape→Gather→Unsqueeze、嵌套 Reshape、训练专用节点。onnxsim 通过「实际跑一遍 + 常量折叠」把它们大幅精简:

bash 复制代码
# 命令行
python -m onnxsim input.onnx output_sim.onnx

# 注意:带动态维度时建议传 --overwrite-input-shape 指定一个具体 shape 来折叠
python -m onnxsim bert_dyn.onnx bert_sim.onnx \
    --overwrite-input-shape 1,16 1,16 1,16 \
    --dynamic-input-shape
python 复制代码
# Python API
from onnxsim import simplify
import onnx

model = onnx.load("resnet50.onnx")
model_sim, check = simplify(model)
assert check, "化简失败"
onnx.save(model_sim, "resnet50_sim.onnx")

化简后通常:节点数变少、图更干净、某些引擎兼容性更好、速度略有提升。建议作为量化和上 TensorRT 之前的标准预处理,但化简后要重新做数值对齐。

12.2 量化:FP32 → INT8

量化把权重/激活从 32 位浮点压到 8 位整数,模型体积约缩到 1/4,推理在支持 INT8 的硬件上可提速 2~4 倍。ONNX Runtime 提供了完整的量化工具链。

① 动态量化(最简单,适合权重主导的小模型 / Linear 多的模型)

python 复制代码
from onnxruntime.quantization import quantize_dynamic, QuantType

quantize_dynamic(
    model_input="resnet50_sim.onnx",
    model_output="resnet50_int8.onnx",
    weight_type=QuantType.QInt8,   # 权重量化成 INT8
)

动态量化在运行时根据观察到的数据范围即时量化激活,无需校准集,开箱即用。对全连接/MatMul 占比高的模型(BERT 等)效果不错。

② 静态量化(效果更好,需要校准数据集)

静态量化提前用一批有代表性的数据统计激活范围,运行时不再动态计算,速度更快、精度通常更可控,是 CNN 上线的主力:

python 复制代码
from onnxruntime.quantization import (
    quantize_static, CalibrationDataReader, QuantType, QuantFormat,
)
import numpy as np
import onnxruntime as ort

class MyCalibrationReader(CalibrationDataReader):
    def __init__(self, samples, input_name):
        self.input_name = input_name
        # samples: 若干真实预处理后的 numpy 输入
        self.data = iter([{input_name: s} for s in samples])

    def get_next(self):
        return next(self.data, None)

sess = ort.InferenceSession("resnet50_sim.onnx", providers=["CPUExecutionProvider"])
input_name = sess.get_inputs()[0].name

# 用真实业务数据做校准(几十~几百条有代表性的样本)
calib_samples = [np.random.randn(1, 3, 224, 224).astype(np.float32) for _ in range(100)]
reader = MyCalibrationReader(calib_samples, input_name)

quantize_static(
    model_input="resnet50_sim.onnx",
    model_output="resnet50_static_int8.onnx",
    calibration_data_reader=reader,
    quant_format=QuantFormat.QDQ,          # QDQ:插入 QuantizeLinear/DequantizeLinear,兼容性/精度更好
    weight_type=QuantType.QInt8,
    activation_type=QuantType.QUInt8,
)

QOperator vs QDQ 怎么选?

  • QDQ(QuantizeLinear/DequantizeLinear) :显式标出量化点,图更直观,TensorRT 等硬件后端识别友好,现代流程首选;
  • QOperator :直接使用量化版本算子(如 ConvInteger),较老的方案。

③ 量化掉点了怎么办(进阶手段)

  • 混合精度 :对敏感层(通常是第一层、最后一层、注意力/Softmax 附近)关闭量化、保留高精度;可用 quantize_static(..., op_types_to_quantize=...) 或节点排除列表精细控制;
  • 更好的校准方法 :默认用最小/最大(MinMax),可换成基于熵的 Percentile / Entropy 校准,寻找更优截断阈值;
  • 量化感知训练(QAT):在训练时模拟量化误差,精度最稳,是掉点严重时的终极方案。ORT/各自训练框架支持伪量化节点训练后再导出。

量化后务必回到第 5.3 节:不仅看数值误差,更要看真实任务指标(分类准确率、召回、mAP)是否在可接受范围。


13. 大模型(LLM / Transformer)的 ONNX 导出

前面讲的 BERT 还属于「一次前向、静态序列」的范畴。到了真正的 LLM,故事完全不同------核心挑战是 KV Cache 与自回归解码。这一节是很多教程跳过、但生产中最关键的部分。

13.1 为什么 LLM 的导出是另一套玩法

LLM 推理是自回归 的:生成第 t t t 个 token,要用到前 t − 1 t-1 t−1 个 token 的注意力结果。如果每步都重算全部前缀,复杂度是 O ( n 2 ) O(n^2) O(n2) 且随生成长度爆炸。KV Cache 的作用就是把历史 token 的 Key/Value 缓存下来,每步只算新 token 的 Q,与缓存的 K/V 做注意力------单步复杂度降到 O ( n ) O(n) O(n)。

这意味着导出时模型必须有两组输入:

复制代码
增量输入: input_ids (batch, 1)           ← 每步只喂 1 个新 token
          attention_mask (batch, seq)
          position_ids (batch, 1)         ← 关键!要告诉模型当前是新序列还是缓存
缓存输入: past_key_values.0.key (batch, num_heads, past_len, head_dim)
          past_key_values.0.value (...)
          ... (每一层都有一组)
缓存输出: present.0.key / present.0.value ...  ← 更新后的缓存,要喂回下一步

核心闭环 :past_*(入)→ 模型 → present_*(出)→ 下一轮把 present_* 作为 past_* 喂回去。

13.2 HuggingFace 官方的 LLM 导出方案

HuggingFace 的 optimum 库对主流 LLM(LLaMA、Qwen、OPT、GPT 系)封装好了 ONNX 导出,是首选方案:

bash 复制代码
# 安装
pip install optimum[onnxruntime]

# 导出 LLaMA 风格模型为 ONNX(自动处理 KV Cache、动态轴)
optimum-cli export onnx \
    --model meta-llama/Llama-3.2-1B \
    --task text-generation-with-past \
    --opset 17 \
    llama_onnx/

# --task text-generation-with-past:关键,带 KV Cache
# 会在 llama_onnx/ 下生成 decoder_model_merged.onnx + 若干外部数据文件

导出的模型通常是一个合并的 decoder(含 KV Cache 的 present/past 输入输出),而非简单的单文件。

13.3 手动导出的关键:动态轴与 KV Cache 声明

如果你想自己用 torch.onnx.export 导出(例如自定义模型),KV Cache 的动态轴声明是难点。给一个结构示意:

python 复制代码
import torch
from transformers import AutoModelForCausalLM, AutoConfig

model_id = "gpt2"  # 换成任意 causal LM
model = AutoModelForCausalLM.from_pretrained(model_id).eval()
config = AutoConfig.from_pretrained(model_id)

num_layers  = config.n_layer
num_heads   = config.n_head
head_dim    = config.n_embd // config.n_head

# 构造带 KV Cache 的样例输入
batch, past_len = 1, 8
input_ids      = torch.ones(batch, 1, dtype=torch.long)          # 只喂 1 个新 token
attention_mask = torch.ones(batch, past_len + 1, dtype=torch.long)
position_ids   = torch.arange(past_len, past_len + 1).unsqueeze(0)

# 每家实现的 past 布局不同,以 (batch, num_heads, past_len, head_dim) 为例
past = []
for _ in range(num_layers):
    k = torch.randn(batch, num_heads, past_len, head_dim)
    v = torch.randn(batch, num_heads, past_len, head_dim)
    past.append((k, v))          # HF v4.x 用扁平 tuple:(k0,v0,k1,v1,...)

# 组装输入顺序(以 HF 约定为准:input_ids, past..., attention_mask, position_ids)
inputs = (input_ids, *[t for kv in past for t in kv], attention_mask, position_ids)

dynamic_axes = {
    "input_ids":      {0: "batch", 1: "one"},
    "attention_mask": {0: "batch", 1: "past_plus_one"},
    "position_ids":   {0: "batch", 1: "one"},
    "logits":         {0: "batch", 1: "one"},
}
for i in range(num_layers):
    dynamic_axes[f"past_key_values.{i}.key"]   = {0: "batch", 2: "past"}
    dynamic_axes[f"past_key_values.{i}.value"] = {0: "batch", 2: "past"}
    dynamic_axes[f"present.{i}.key"]           = {0: "batch", 2: "past_plus_one"}
    dynamic_axes[f"present.{i}.value"]         = {0: "batch", 2: "past_plus_one"}

# 注意:LLM 往往 >2GB,必须用外部数据格式!
torch.onnx.export(
    model,
    inputs,
    "llm/model.onnx",
    opset_version=17,
    input_names=["input_ids", *[f"past_key_values.{i}.{n}"
                                for i in range(num_layers) for n in ("key", "value")],
                 "attention_mask", "position_ids"],
    output_names=["logits", *[f"present.{i}.{n}"
                              for i in range(num_layers) for n in ("key", "value")]],
    dynamic_axes=dynamic_axes,
    use_external_data_format=True,     # ★ 大模型必备
)

⚠️ 强烈建议 :除非你在做研究或特殊定制,直接用 optimum-cli 或各家官方导出脚本 。手动导出 LLM 的 KV Cache 极容易因为 past 布局差异(HF 不同版本从 past_key_values 的嵌套 tuple 到扁平 tuple 改过好几次)而出错。手动导出更适合理解原理。

13.4 生成循环:把 ONNX 当 KV Cache 引擎用

导出后,你要自己写解码循环,把握住「present 喂回 past」这个闭环:

python 复制代码
import numpy as np
import onnxruntime as ort

sess = ort.InferenceSession("llm/model.onnx", providers=["CPUExecutionProvider"])
print("输入:", [i.name for i in sess.get_inputs()][:4], "...")
print("输出:", [o.name for o in sess.get_outputs()][:4], "...")


class KVCacheRunner:
    def __init__(self, session, num_layers):
        self.sess = session
        self.num_layers = num_layers
        self.past = None

    def step(self, input_ids, attention_mask, position_ids):
        feeds = {
            "input_ids": input_ids,
            "attention_mask": attention_mask,
            "position_ids": position_ids,
        }
        # 第一轮没有缓存时,喂全零的 past;后续轮喂上一轮的 present
        if self.past is not None:
            for i in range(self.num_layers):
                feeds[f"past_key_values.{i}.key"]   = self.past[f"present.{i}.key"]
                feeds[f"past_key_values.{i}.value"] = self.past[f"present.{i}.value"]

        out_names = ["logits"] + [f"present.{i}.{n}"
                                  for i in range(self.num_layers)
                                  for n in ("key", "value")]
        outs = self.sess.run(out_names, feeds)
        self.past = {"logits": outs[0], **{n: o for n, o in zip(out_names[1:], outs[1:])}}
        return outs[0]   # logits: (batch, 1, vocab)


# 伪代码:贪心解码
runner = KVCacheRunner(sess, num_layers=12)
tokens = [tokenizer.bos_token_id]
for t in range(max_new_tokens):
    ids = np.array([[tokens[-1]]], dtype=np.int64)
    pos = np.array([[t]], dtype=np.int64)
    mask = np.ones((1, len(tokens)), dtype=np.int64)
    logits = runner.step(ids, mask, pos)          # (1, 1, vocab)
    next_id = int(logits[0, -1].argmax())
    if next_id == tokenizer.eos_token_id:
        break
    tokens.append(next_id)
print(tokenizer.decode(tokens))

几个 LLM 导出的工程要点:

要点 说明
必须用外部数据格式 LLM 体积普遍 >2GB,撞 protobuf 上限;optimum 会自动处理
past 布局随版本变 HF past_key_values 结构多次变更,导出前先确认当前版本约定
position_ids 必须动态 增量解码时它表示「当前绝对位置」,写死会导致位置编码错误
算子兼容性 MultiHeadAttention、RotaryEmbedding(RoPE)等在新 opset 才友好;老引擎需分解成基础算子
量化对 LLM 收益巨大 INT8 常掉点明显,权重 W4A16(如 INT4) 才是 LLM 量化的主流方向
别只看单步延迟 要分别报告 Prefill(首 token,算力密集) 与 Decode(后续 token,访存密集) 延迟

📌 一个诚实提醒 :ONNX Runtime 能跑 LLM,但对大模型的极致优化(PagedAttention、Continuous Batching、投机解码)上,专用推理框架如 vLLM、TensorRT-LLM、SGLang 更成熟。ONNX 在 LLM 场景的价值更多在于中小模型的端侧/嵌入式部署 、跨硬件的统一分发 ,以及作为 TensorRT-LLM 等引擎的前置格式。选型时别被"什么都能跑"迷惑。


14. 走向生产:ONNX 的推理生态地图

ONNX 最大的价值是它背后的「后端全家桶」。转成 ONNX 后,按目标硬件选择引擎:

目标硬件 / 场景 推荐引擎 说明
跨平台 CPU / 通用服务 ONNX Runtime 最通用,Win/Linux/Mac、Python/C++/C#/Java 全支持
NVIDIA 数据中心 GPU TensorRT 极致延迟/吞吐,支持 FP16/INT8,ONNX 是其首选输入格式
Intel CPU / 核显 OpenVINO Intel 硬件上 INT8 推理很强
Windows / 跨平台 GPU ONNX Runtime + CUDA/DirectML DirectML 可兼容各类 Windows GPU
iPhone / iPad Core ML 苹果生态,通过 coremltools 从 ONNX/其他格式转
Android NNAPI / ORT NNAPI EP 调用手机 NPU/GPU
浏览器 ORT Web / onnxruntime-web WebAssembly + WebGPU,纯前端推理

14.1 二选一:ONNX Runtime 还是 TensorRT?

这是 GPU 部署里最常见的选择题,给一张直接的对照:

维度 ONNX Runtime (+CUDA EP) TensorRT
上手成本 低,pip 即用 高,需要编译 engine、调 profile
极致性能 良好 最好(kernel 自动调优、层融合更激进)
动态 shape 天然支持 需要预先声明 min/opt/max 三档
部署形态 一个 .onnx 文件跨平台 每台机器/每种 shape 可能需重新 build engine
适合场景 快速上线、多硬件分发、CPU/GPU 混合 固定硬件、固定 shape、榨性能

经验结论 :先用 ORT 上线拿基线,等性能成为瓶颈、且硬件固定时,再切 TensorRT 抠最后那 2~3 倍。 不要一上来就迷信 TensorRT。

14.2 端侧与浏览器:ONNX 的另一片战场

除了服务器,ONNX 在端侧和 Web 上同样重要,这是很多教程忽略的场景:

① 浏览器内推理(ORT Web)------纯前端、无需后端、数据不出本地:

html 复制代码
<script src="https://cdn.jsdelivr.net/npm/onnxruntime-web/dist/ort.min.js"></script>
<script>
  async function run() {
    // 指定用 WebGPU / WASM 后端
    const session = await ort.InferenceSession.create(
      "resnet50.onnx",
      { executionProviders: ["webgpu", "wasm"] }
    );
    const data = new Float32Array(1 * 3 * 224 * 224);
    const tensor = new ort.Tensor("float32", data, [1, 3, 224, 224]);
    const outputs = await session.run({ input: tensor });
    console.log(outputs.output.dims);
  }
  run();
</script>

② 手机 / 嵌入式 :ORT 提供 iOS / Android 的移动端包,也可通过 coremltools 转 Core ML、通过转换器转各家 NPU 的格式。

③ 边缘设备(树莓派、Jetson、国产盒子) :小模型 + INT8 量化后,ORT 能直接在 ARM CPU 上跑;Jetson 则有 TensorRT。端侧部署的核心永远是「量化 + 小模型 + 固定 shape」。

💡 一个端侧经验 :给端侧导出模型时,优先固定 shape、优先 INT8、优先精简图 。端侧对内存和功耗极其敏感,动态维度和冗余算子在这里都是负担。同一套模型,服务器版和端侧版往往要用不同参数导出两次。

14.3 ONNX 与其它格式的关系

开发中你会遇到一堆模型格式,别混淆它们各自的定位:

格式 定位 与 ONNX 的关系
ONNX 跨框架交换格式(中间层) 本文主角
.pt / .pth PyTorch 原生权重 上游,需导出成 ONNX
SavedModel / .h5 TensorFlow 原生 上游,需 tf2onnx 转
.engine / .plan TensorRT 编译产物 下游,由 ONNX 编译而来
IR(.xml/.bin) OpenVINO 格式 下游,由 ONNX 转换而来
.mlmodel / .mlpackage Core ML 格式 下游,由 ONNX/其它格式转
GGUF llama.cpp 的 LLM 格式 平行生态,LLM 端侧常用,非 ONNX
safetensors 权重安全存储格式 平行的"存储"格式,非"计算图"格式

记住一句话 :ONNX 是计算图 + 权重的交换中间层,上游是训练框架,下游是推理引擎。它不是终点,而是「枢纽」。

以最常见的「ONNX → TensorRT」为例 ,新版 TensorRT 可以直接吃 ONNX(背后用官方的 trtexec / ONNX GraphSurgeon):

bash 复制代码
# FP16
trtexec --onnx=resnet50_sim.onnx \
        --saveEngine=resnet50_fp16.engine \
        --fp16

# INT8(需要校准缓存或显式精度)
trtexec --onnx=resnet50_sim.onnx \
        --saveEngine=resnet50_int8.engine \
        --int8 \
        --calib=calibration.cache

# 动态 shape:给出 min / opt / max 三档
trtexec --onnx=bert_sim.onnx \
        --minShapes=input_ids:1x8,attention_mask:1x8,token_type_ids:1x8 \
        --optShapes=input_ids:8x64,attention_mask:8x64,token_type_ids:8x64 \
        --maxShapes=input_ids:32x256,attention_mask:32x256,token_type_ids:32x256 \
        --fp16 --saveEngine=bert.engine

trtexec 还会直接打印各层耗时和总的吞吐/延迟,是调优的好帮手。

一条我常用的生产流水线长这样:

复制代码
PyTorch/TF 训练
      │  导出器
      ▼
 原始 ONNX
      │  onnxsim 化简 + shape_inference
      ▼
 精简 ONNX  ──数值对齐──┐
      │                 │
      ├─ CPU:  ORT(静态/动态量化 INT8)
      ├─ GPU:  TensorRT(FP16/INT8 + profile)
      ├─ Intel:OpenVINO(INT8)
      └─ 端侧: Core ML / NNAPI

把每一步的产物和数值校验固化成脚本/CI,模型迭代时就能一键重放整条部署链路,这是工程化成熟度的标志。


15. 进阶:当你遇到不支持的算子

大多数情况下,第 11 节「用标准算子改写 / 换导出器」已经够用。这里简单介绍两种高级手段,让你知道边界在哪。

15.1 为老导出器注册自定义 symbolic

如果某个算子(包括你自己的 autograd.Function)实在无法改写,可以告诉导出器「这个算子该翻译成什么 ONNX 子图」:

python 复制代码
import torch
from torch.onnx import register_custom_op_symbolic
from torch.onnx.symbolic_helper import parse_args

@parse_args("v", "i")
def my_op_symbolic(g, x, axis):
    # 用 ONNX 算子拼出等价实现,例如这里退化成一个 ReduceMean
    return g.op("ReduceMean", x, axes_i=[axis], keepdims_i=1)

# 注册到对应算子名和 opset
register_custom_op_symbolic("mynamespace::my_op", my_op_symbolic, opset_version=17)

这要求你对 ONNX 算子集比较熟,属于「最后手段」。

15.2 用 onnxscript 写自定义函数

onnxscript 让你用一种「贴近 Python 的前端语言」直接写 ONNX 子图,是 TorchDynamo 新导出器处理复杂/自定义算子的基础,也是官方主推方向:

python 复制代码
from onnxscript import script
from onnxscript import opset17 as op

@script()
def MyAddMean(x, y):
    s = op.Add(x, y)
    return op.ReduceMean(s, axes=[-1], keepdims=1)

写出的函数既能当普通 ONNX 子图内联,也能在支持的引擎上作为自定义函数调用。长期看,复杂模型的算子问题优先在 onnxscript + Dynamo 这条新链路上解决。


16. 生产部署:走出 Python,用 C++ 跑推理

Python 适合验证和快速原型,但线上服务、桌面软件、嵌入式 通常要用 C++/C#/Java 的 ORT API------不仅性能开销更小,还能摆脱 Python 运行时和 GIL 的束缚。这里给一个 C++ 的 ORT 完整推理骨架,是生产集成的常见起点。

cpp 复制代码
// onnx_infer.cpp  ------ 编译:g++ onnx_infer.cpp -lonnxruntime -std=c++17
#include <onnxruntime_cxx_api.h>
#include <vector>
#include <iostream>
#include <array>

int main() {
    // 1) 环境与 Session
    Ort::Env env(ORT_LOGGING_LEVEL_WARNING, "onnx_demo");
    Ort::SessionOptions opts;
    opts.SetIntraOpNumThreads(4);
    opts.SetGraphOptimizationLevel(GraphOptimizationLevel::ORT_ENABLE_ALL);

    Ort::Session session(env, "resnet50.onnx", opts);

    // 2) 输入输出名(用 UTF-8 字符串,须保证生命周期)
    const char* input_names[]  = {"input"};
    const char* output_names[] = {"output"};

    // 3) 构造输入张量:NCHW = (1,3,224,224)
    std::array<int64_t, 4> shape{1, 3, 224, 224};
    std::vector<float> input_data(1 * 3 * 224 * 224, 0.0f);   // 填充你的预处理结果

    Ort::MemoryInfo mem = Ort::MemoryInfo::CreateCpu(
        OrtArenaAllocator, OrtMemTypeDefault);
    Ort::Value input_tensor = Ort::Value::CreateTensor<float>(
        mem, input_data.data(), input_data.size(), shape.data(), shape.size());

    // 4) 推理
    auto outputs = session.Run(
        Ort::RunOptions{nullptr},
        input_names, &input_tensor, 1,
        output_names, 1);

    // 5) 取结果
    float* out = outputs[0].GetTensorMutableData<float>();
    auto out_shape = outputs[0].GetTensorTypeAndShapeInfo().GetShape();
    int64_t num_classes = out_shape[1];

    // argmax 得到预测类别
    int64_t best = 0; float best_val = out[0];
    for (int64_t i = 1; i < num_classes; ++i) {
        if (out[i] > best_val) { best_val = out[i]; best = i; }
    }
    std::cout << "预测类别: " << best
              << "  置信度: " << best_val << std::endl;
    return 0;
}

C++ 侧的几个工程要点:

要点 说明
张量命名生命周期 Ort::Value 引用外部内存时,源数据必须在 Run 期间存活;别用临时 vector
复用 Session Session 创建开销大,且内含已优化的图------全局复用,别每次请求都建
线程安全 Session 的 Run 是线程安全的,可以多线程并发调用同一个 Session
动态 shape C++ 侧用 GetInputNameAllocated 拿名字,按实际 shape 构造输入张量
大模型外部数据 加载 model.onnx 时,同目录的 .data 文件必须存在
内存池 高并发场景配置 OrtArenaAllocator,减少分配抖动

🎯 一个完整的生产服务架构通常是:C++ 侧负责「预处理 → ORT 推理 → 后处理」的热路径,Python 只在离线侧负责模型转换与校验。把预处理也下沉到 C++(避免 Python 来回转 numpy),往往还能再带来可观收益。


17. 生产级转换脚本(可直接复用)

最后,把前面所有最佳实践串成一个带校验的脚本,支持动态 batch、化简、可选量化,并在每一步做数值对齐:

python 复制代码
"""
export_onnx.py ------ 生产级 PyTorch → ONNX 转换模板
用法:
    python export_onnx.py
"""
import numpy as np
import torch
import torch.nn as nn
import onnx
import onnxruntime as ort
from onnxsim import simplify


def export_to_onnx(
    model: nn.Module,
    dummy_inputs,                 # tuple[torch.Tensor]
    onnx_path: str,
    input_names,
    output_names,
    dynamic_axes,
    opset: int = 17,
    do_simplify: bool = True,
    rtol: float = 1e-3,
    atol: float = 1e-5,
):
    model.eval()

    # ---------- 1. 导出 ----------
    with torch.no_grad():
        torch.onnx.export(
            model,
            tuple(dummy_inputs),
            onnx_path,
            export_params=True,
            opset_version=opset,
            do_constant_folding=True,
            input_names=input_names,
            output_names=output_names,
            dynamic_axes=dynamic_axes,
        )
    print(f"[1/4] 导出完成 -> {onnx_path}")

    # ---------- 2. 结构校验 ----------
    onnx_model = onnx.load(onnx_path)
    onnx.checker.check_model(onnx_model)
    print("[2/4] 结构校验通过")

    # ---------- 3. 化简 ----------
    if do_simplify:
        onnx_model, ok = simplify(
            onnx_model,
            overwrite_input_shapes={
                name: list(t.shape) for name, t in zip(input_names, dummy_inputs)
            },
        )
        assert ok, "onnxsim 化简失败"
        onnx.save(onnx_model, onnx_path)
        print("[3/4] 图化简完成")
    else:
        print("[3/4] 跳过化简")

    # ---------- 4. 数值对齐 ----------
    with torch.no_grad():
        ref = model(*dummy_inputs)
    ref_outs = ref if isinstance(ref, (tuple, list)) else (ref,)

    sess = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
    feeds = {name: t.detach().cpu().numpy()
             for name, t in zip(input_names, dummy_inputs)}
    ort_outs = sess.run(None, feeds)

    for i, (r, o) in enumerate(zip(ref_outs, ort_outs)):
        np.testing.assert_allclose(r.cpu().numpy(), o, rtol=rtol, atol=atol)
        print(f"[4/4] 输出[{i}] 数值一致,max abs diff = "
              f"{np.abs(r.cpu().numpy() - o).max():.3e}")

    print("✅ 全部检查通过,模型可用于部署")
    return onnx_path


if __name__ == "__main__":
    import torchvision

    model = torchvision.models.resnet50(
        weights=torchvision.models.ResNet50_Weights.DEFAULT
    )
    dummy = (torch.randn(2, 3, 224, 224),)

    export_to_onnx(
        model=model,
        dummy_inputs=dummy,
        onnx_path="resnet50_prod.onnx",
        input_names=["input"],
        output_names=["output"],
        dynamic_axes={
            "input":  {0: "batch"},
            "output": {0: "batch"},
        },
        opset=17,
    )

大模型(>2GB)在这个脚本里加 use_external_data_format=True;需要量化就在数值对齐后接第 12 节的量化函数,并在量化后再做一次任务级指标验证。


18. 报错速查表

遇到问题先查这张表,按图索骥:

报错 / 现象关键词 大概率原因 快速处理
每次推理结果都不一样 没切 eval(),BN/Dropout 在训练态 model.eval()
operator XXX is not supported 算子无 ONNX symbolic 升 opset / 改写算子 / 用 Dynamo
IR version ... not supported IR/opset 超过引擎支持 降低 opset_version
喂别的 batch 就 shape 报错 维度写死 配置 dynamic_axes
输入列表里有一堆权重名 initializer 被当成输入 清理 graph.input(坑 5)
大模型导出超限/打不开 撞上 protobuf 2GB 上限 use_external_data_format
Unexpected input data type dtype 不匹配 对齐 get_inputs() 的 dtype
装了 GPU 版却很慢 算子回退到 CPU EP 查日志看 EP 分配,升级/换算子
Softmax/归一化结果错乱 axis / NCHW·NHWC 错位 Netron 逐节点核对
量化后掉点严重 INT8 精度损失 静态校准 / 敏感层混精 / QAT
结果对但速度偏慢 残留大量 Transpose / 冗余节点 --nchw、onnxsim 化简
控制流分支丢失 trace 只录了样例分支 用 Dynamo 或消除数据依赖分支
转 TensorRT 报算子不认识 图里混入了 com.microsoft 私有算子 查 Netron 的 domain 字段,换导出路径
BN 模型有稳定小偏差 epsilon 不一致 核对 epsilon 属性
变长 batch 结果不一致 attention_mask 漏传/错传 显式导出并喂真实 mask
LLM 加载失败 >2GB 未用外部数据 use_external_data_format=True
检测框位置整体偏移 letterbox 的 padding 未在反算时还原 用 dw/dh/ratio 反算(9.3)

19. 总结与展望

回到开头那个「下周上线」的场景,现在你应该有了完整的武器库:

  1. ONNX 是深度学习世界的中间表示,用 Protobuf 描述计算图、算子和权重,解耦了训练框架与推理硬件;
  2. 转换只是第一步,专业流程是「导出 → 结构校验 → 化简 → 数值对齐 → 量化/编译 → 任务级验证」,缺一不可;
  3. 动态维度、多输入输出、dict 输出包装、外部数据格式是处理真实模型的四个高频技巧;
  4. opset 是兼容性的生命线 :云端可上 17,端侧保守到 13,且要警惕 com.microsoft 私有算子;
  5. PyTorch 2.x 的 TorchDynamo + onnxscript 正在成为新的默认,复杂模型优先在新链路上解决算子问题;
  6. 性能调优有方法论:benchmark(P50/P99)→ profiling(找慢节点)→ 线程/IO Binding/量化逐层优化,而不是盲猜;
  7. 大模型是新的深水区:KV Cache 的 past/present 闭环、外部数据格式、Prefill/Decode 分离,都是必修课;
  8. ONNX 的真正威力在其后端生态------ORT、TensorRT、OpenVINO、Core ML、ORT Web 让「一次导出,到处高性能运行」成为现实。

最后给一条务实的落地路线,如果你正要开始一个部署项目:

复制代码
第 1 天:跑通最小闭环
  导出 ResNet/BERT → check_model → 数值对齐 → ORT 推理

第 1 周:把生产要素补齐
  动态 shape、多输入、外部数据格式、onnxsim 化简、基准测试

第 2 周:性能与精度
  Profiling 找瓶颈 → 线程/IO Binding 调优 → INT8 量化 + 任务指标验证

第 3 周:走向目标硬件
  GPU→TensorRT / Intel→OpenVINO / 端侧→量化+固定 shape;C++ 集成

展望未来,几个趋势值得关注:

  • 大模型(LLM / 多模态)成为 ONNX 的新战场:超长动态序列、KV Cache、投机解码等给图表达和引擎带来新挑战,ONNX 规范和 ORT 也在快速跟进;但极致性能仍要看 TensorRT-LLM / vLLM 等专用引擎;
  • 量化持续下沉:INT8 已普及,INT4 / NF4 等低位宽格式与权重只量化(W4A16)正在端侧和大模型推理中普及;
  • 导出器与编译器进一步融合:Dynamo → onnxscript → 各硬件后端的一体化链路会越来越顺滑,「部署」这件事的门槛会持续降低;
  • 统一编译栈的竞争:ONNX 之外,MLIR、TVM、Apache TVM 等编译栈也在演进,但 ONNX 凭借生态和厂商支持,仍是「交换格式」这一层最稳的赢家。

如果这篇文章帮你少加了一个晚班,欢迎点赞、在看、转发三连 🙌

我会持续输出 模型部署、推理优化、大模型工程化 方向的硬核内容------既有可以直接抄走的代码,也有背后的原理和踩坑实录。关注我,下一篇我们聊聊 TensorRT 极致优化 / 大模型的 ONNX 导出与 KV Cache 加速。

你的每一次关注,都是我继续写下去的动力。 我们下篇见 👋


相关推荐
zh路西法13 小时前
【四足强化学习经典算法】(一):RMA——盲视运动控制与快速适应模块
pytorch·强化学习·1024程序员节·四足机器人·dagger·pyhon·rma
Be for thing20 小时前
PyCharm安装与PyTorch项目配置|Jupyter Notebook与PyCharm使用对比
pytorch·学习·1024程序员节
zh路西法1 天前
【模型式强化学习】(二)PETS——从概率动力学模型到模型预测控制(MPC)
pytorch·python·强化学习·pe·mpc·ces·pets
段一凡-华北理工大学2 天前
高炉炼铁机器视觉与智能识别十八讲~系列文章17:实战案例:一个风口监控系统的完整搭建
人工智能·pytorch·深度学习·机器学习·计算机视觉·高炉风口监测系统
Yolanda_20222 天前
PyTorch 线性层nn.Linear底层数学原理
人工智能·pytorch·python
专注仿真2 天前
多智能体强化学习实战:从 Dec-POMDP 到 QMIX/MADDPG/MAPPO 的可运行 PyTorch 实现
人工智能·pytorch·python
Web3&Basketball2 天前
TorchTPU 迁移实战:3 个 CUDA 卡点
pytorch·性能调优·cuda·tpu·模型迁移·torchtpu
田里的水稻3 天前
FA_CNN卷积神经网络十二_CIFAR10全流程PyTorch
人工智能·pytorch·深度学习·神经网络·cnn
论文复现现场3 天前
RTX 3090 云GPU租赁价格对比:8家平台配置、计费方式与环境测试
linux·运维·人工智能·pytorch·深度学习·云计算·3090
liwulin05063 天前
【PYTORCH】迷你DECODER-ONLY模型
pytorch·python·深度学习