1. ONNX 到底解决什么问题
先抛一个比较直白的类比:ONNX 在大模型和 AI 推理生态里的角色,有点像音频领域的 MP3,或者文档领域的 PDF。你不需要关心它背后是哪个软件做的、在什么系统上跑的,只要格式统一,换个环境依然能"打开"。ONNX 本质上做的就是这件事——把各种深度学习框架训练出来的模型,转换成一个统一的、标准化的计算图描述格式。
这句话听起来很简单,但往深了说,它解决的是 AI 工程化里最头疼的"格式孤岛"难题。
在 ONNX 出现之前,PyTorch 训练出来的模型只能靠 PyTorch 自己推理,TensorFlow 训练的模型也只能在自己生态里跑,Caffe 的模型更是出了名的难迁移。如果业务方用的是 Java 后端,想直接调用一个 PyTorch 模型,要么用 PyTorch Serve 单独起一个 Python 服务,要么花大力气把模型重写一遍。这两种方案在工程上的成本都不低:前者引入了额外的进程通信和运维负担,后者简直是一场灾难——深度学习模型动辄几十上百层,逐层复现网络结构,任何一个参数对不上,输出结果就完全不对。
ONNX 把这个问题拆成了两个阶段:训练阶段继续用 PyTorch 或 TensorFlow,得到模型后导出为 ONNX 格式的中间文件;推理阶段再用 ONNX Runtime 或者转成 TensorRT、OpenVINO 等格式执行。训练和推理解耦之后,整个架构的灵活性一下就上来了——模型训练团队不用关心线上部署环境是 Java 还是 C++,部署团队也不用关心模型是用什么框架训出来的。
大模型时代到来之后,ONNX 的存在感不但没有减弱,反而变得更重要了。现在做 LLM 推理,大家满脑子都是 vLLM、TensorRT-LLM、Ollama 这些名字,但它们内部其实都很重视 ONNX 这条链路。原因在于大模型的部署形态已经变得极其复杂:既要跑在云端 A100/H100 上,也要跑在国产加速卡上,还要兼顾 PC 端的 CPU 环境甚至移动端。一个统一的中间表示,能帮你省下大量重复适配的功夫。
这篇内容我会从 ONNX 的核心设计讲起,然后带你完整走一遍模型转换到推理的实操流程,最后聊一聊在大模型和各类推理引擎并存的局面下,ONNX 到底应该放在什么位置。无论你是刚接触模型部署的新手,还是正在折腾模型迁移的老手,应该都能从中找到点有用的东西。
2. ONNX 的核心设计:计算图、算子与格式细节
2.1 ONNX 模型文件里到底装了什么
拿一个 .onnx 文件在编辑器里打开(或者用 Netron 可视化),你会看到它本质是一个 protobuf 序列化的数据文件。里面的结构分几大块:计算图(Graph)、算子集合(Opset)、模型元信息,以及可选的权重参数。
计算图是 ONNX 的核心。它描述的是数据从输入到输出的流动过程:图中的每个节点(Node)代表一个算子操作,比如 Conv、MatMul、Add 这些;边(Edge)代表张量(Tensor)在算子之间如何传递。每个节点都定义了输入张量的名称、输出张量的名称,以及算子自身的属性参数。整个图结构可以看成一张有向无环图(DAG),数据从输入节点流入,经过一层层算子变换,最终从输出节点流出。
这里要特别强调一下,ONNX 存的不是训练好的"神经网络对象",而是一个逐个算子展开的静态计算流程。PyTorch 里的 nn.Module 是面向训练设计的,里面包含了很多训练特有的逻辑(比如 dropout 的随机掩码、BN 层的滑动均值更新),而 ONNX 只关心推理路径上的算子序列。这也是为什么导出时经常要指定 opset_version——不同版本的算子集合支持的能力不同,新版本通常会增加新算子或修改旧算子的行为描述。
权重参数也直接打包在 .onnx 文件里。PyTorch 模型的 state_dict 是分散的参数集合,但 ONNX 会按照计算图的拓扑顺序把权重组织起来,作为图的初始化器(Initializer)内嵌在文件里。这也是 .onnx 文件通常比原模型更大的原因之一,但好处是不用额外管理权重文件,一个文件拿过去就能直接跑推理。
2.2 算子集合(Opset)和 IR 版本为什么重要
理解 Opset 是 ONNX 绕不开的一关。ONNX 的算子集合是版本化的,每个版本的算子集合新增或修改了一些算子定义。比如 opset 11 引入了一些动态 shape 相关的修改,opset 13 优化了 Reduce 系列算子的行为,opset 17 之后对字符串和复数类型的支持更完善。导出模型时指定的 opset_version 决定了模型使用的是哪个版本的语义。
这个问题在实践中非常关键,因为它直接影响模型的可移植性。如果你在 opset 11 下导出模型,遇到一个 opset 13 才支持的新算子,就得想办法绕过或者升级 opset 版本。反过来,如果你用了很新的 opset,但推理端的 ONNX Runtime 版本较老,可能也会因为不支持新的算子集合而报错。
IR 版本(Intermediate Representation version)是另一个容易忽略的配置。它定义了模型文件本身的格式规范,比如数据类型的表示方式、图结构的序列化方式等。IR 版本和 opset 版本需要匹配,导出工具通常会帮你处理好,但如果你手工修改过模型文件或者跨版本转换过,就要多留意一下这两者的兼容性。
实践中我的建议是:除非有特殊算子需求,否则不要追求过新的 opset。尽量选择一个主流推理框架都已经兼容的版本,比如 11、13、17 这几年用的都比较多。太激进地使用新版本 op set,往往会在国产加速卡或者边缘设备的推理引擎上踩到兼容性坑。
2.3 Netron:看 ONNX 模型的必备工具
接触 ONNX 之后,Netron 应该是使用频率最高的可视化工具了,没有之一。它是一个开源的神经网络模型可视化工具,支持 ONNX、TensorFlow Lite、Keras、CoreML 等多种格式。把 .onnx 文件拖进去,就能看到完整的计算图结构、每个节点的输入输出 shape、权重参数的维度,甚至能直接查看某个节点的属性值。
调试模型导出问题时,Netron 的价值特别大。比如你发现导出的模型推理结果不对,先用 Netron 打开看一遍构图是否符合预期,能快速定位是不是某个算子的连接关系搞错了,或者权重有没有对调。还有一些情况下,PyTorch 里看起来正常的操作,导出后可能会被拆成多个细粒度的算子组合(因为 ONNX 的算子粒度往往比 PyTorch 的层粒度更细),Netron 里一眼就能看出来这些组合是否合理。
Netron 还有网页版和桌面版,网页版直接上传本地文件就能用,桌面版支持更大的模型文件。大模型场景下,动辄几个 GB 的 .onnx 文件在网页版里加载会比较吃力,建议直接用桌面版会更顺手一些。
3. 实操一把:把 PyTorch 模型转换成 ONNX
3.1 准备环境和依赖
开始之前,先把环境准备好。最基础的三件套是 PyTorch、ONNX 和 ONNX Runtime。如果后面要可视化检查,可以再装一个 netron 的 Python 包,这样可以直接在代码里启动可视化界面。
我建议用 Python 3.8 以上的环境(3.10、3.11 都没问题),PyTorch 装 2.x 版本,ONNX 装 1.14 以上,ONNX Runtime 装最新稳定版。如果涉及到量化,后面还需要 onnxruntime 的扩展包,这个先按下不表,后面讲量化的时候具体说。
import torch import torch.nn as nn # 定义一个简单的 CNN 模型用于演示 class SimpleCNN(nn.Module): def __init__(self): super().__init__() self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1) self.conv2 = nn.Conv2d(16, 32, kernel_size=3, padding=1) self.fc = nn.Linear(32 * 8 * 8, 10) def forward(self, x): x = torch.relu(self.conv1(x)) x = torch.max_pool2d(x, 2) x = torch.relu(self.conv2(x)) x = torch.max_pool2d(x, 2) x = x.view(x.size(0), -1) x = self.fc(x) return x这里要注意,PyTorch 从 2.0 开始默认是动态图模式,导出 ONNX 时需要把模型切换到推理模式,并且用 torch.no_grad() 包裹推理过程。另外,模型的 forward 方法里如果有 if 分支或者 Python 层的动态逻辑,导出时也要格外小心,ONNX 只能捕捉到实际执行过的路径。
3.2 导出 ONNX 的关键参数设置
PyTorch 导出 ONNX 的核心接口是 torch.onnx.export。这个方法看起来很简单,但参数细节决定了导出成败。
model = SimpleCNN() model.eval() # 切换到推理模式 dummy_input = torch.randn(1, 3, 32, 32) torch.onnx.export( model, dummy_input, "simple_cnn.onnx", export_params=True, opset_version=13, do_constant_folding=True, input_names=["input"], output_names=["output"], dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}} )重点说几个参数背后的逻辑:
export_params=True 表示把权重参数一并导出到 ONNX 文件里。如果你设置成 False,导出的模型就不带权重,只有图结构,这在某些特殊场景(比如只在特定框架里加载结构后手动灌权重)下有用,但常规推理场景请务必保持 True,不然后期用起来会非常痛苦。
do_constant_folding=True 会在导出时对常量子图进行折叠优化。比如某个算子只做了纯常量的数学计算,导出的过程中直接就把结果算好写进文件,推理时就不用再重复计算了。默认是开启的,我建议保持开启,可以让模型更精简,推理速度也有一定提升。
dynamic_axes 是动态维度设置。上面代码里把 batch_size 维度设成了动态,意味着导出的模型在推理时可以接受任意 batch 大小的输入,而不只是固定为 1。这个配置在很多场景下都很有必要,但要注意,动态维度会限制某些推理引擎的优化效果,比如 TensorRT 对动态 batch 的优化通常比静态 shape 差一些。如果你的业务场景里 batch size 是固定的,就不要全部设成动态。
检查一下导出的模型是否正常,可以用 ONNX 自带的校验工具:
import onnx model = onnx.load("simple_cnn.onnx") onnx.checker.check_model(model) print(onnx.helper.printable_graph(model.graph))onnx.checker.check_model 会检查模型结构是否合法,包括图结构是否完整、节点连接是否正确、数据类型的定义是否符合规范等。如果这一步能通过,说明模型本身在格式层面没有问题。
3.3 ONNX Runtime 推理验证
导出完成后,用 ONNX Runtime 跑一遍推理,对比 PyTorch 的输出结果,这是检验转换是否正确最直接的方法。
import onnxruntime as ort import numpy as np # 创建 ONNX Runtime 推理会话 sess_options = ort.SessionOptions() sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL session = ort.InferenceSession("simple_cnn.onnx", sess_options=sess_options) # 准备输入数据 input_data = np.random.randn(1, 3, 32, 32).astype(np.float32) # 获取输入输出名称 input_name = session.get_inputs()[0].name output_name = session.get_outputs()[0].name # 推理 output = session.run([output_name], {input_name: input_data})[0] print("ONNX Runtime 输出 shape:", output.shape) # 同样的输入用 PyTorch 推理对比 with torch.no_grad(): torch_output = model(torch.from_numpy(input_data)).numpy() print("最大误差:", np.max(np.abs(output - torch_output)))这里设置的 graph_optimization_level 就是 ONNX Runtime 的优化等级。ORT_ENABLE_ALL 表示启用全部图优化,包括算子融合、布局优化、常量折叠等。推理引擎在加载模型时会按照这个优化级别对模型进行优化,然后再执行。
最大误差控制在 1e-5 级别以内就说明转换没有问题。如果误差很大,通常有几个可能的原因:一是模型里包含训练特有的层(比如 BN 层在训练和推理模式下行为不同),导出时忘记切到 eval 模式;二是有一些自定义算子没有被 ONNX 支持,导出工具用多个基础算子拼凑代替时引入了精度损失;三是对动态 shape 的处理有偏差。这种误差对比在模型迁移场景下是必做的一步,不要偷懒跳过。
3.4 一个容易踩的坑:控制流的导出
如果你处理的是 LLM 这类大模型,或者任何在 forward 方法里包含 if 条件判断、循环等控制流的模型,导出 ONNX 时就要特别注意了。PyTorch 的 torch.onnx.export 默认走的是 tracing 路线——它只记录实际执行过的算子路径,不会把 if 分支的另一个分支也导出到 ONNX 图里。
举个例子,如果你的模型中有:
def forward(self, x): if x.shape[1] > 10: x = self.layer_a(x) else: x = self.layer_b(x) return x用 dummy_input 导出时,dummy_input 的 shape 决定了哪一个分支会被追踪。如果你设置一个 shape 为 (1, 12, ...) 的 dummy_input,导出的模型里就只有 layer_a 这个分支。线上推理时如果输入变成 (1, 5, ...),模型根本不知道还有 layer_b 这个分支的存在,结果就会出问题。
这种情况下有两种解决思路。一种是通过 torch.onnx.is_onnx_export() 这类条件判断接口,在导出时手动控制走哪个分支;另一种是把静态的 if 判断改成动态的可导算子,让 ONNX 图天然支持不同 shape 的输入。对于大模型来说,很多部署工具(比如 HuggingFace Optimum)已经帮你把这些转换细节处理好了,但如果你是自己手写导出逻辑,这块一定要留意。
4. 不只是转换:ONNX 模型的优化手段
4.1 模型精简:去掉冗余节点和算子融合
拿到一个 ONNX 模型之后,第一件事不是直接上推理引擎,而是先看看模型有没有优化的空间。ONNX 模型本身就有一些优化手段,常见的有这么几类:
第一类是节点精简。深度学习模型在训练框架里定义的时候,为了方便开发者阅读和调试,结构往往比较冗余。比如某个算子做了 A+B,然后又有一个算子做了结果乘 C,导出成 ONNX 后可能是一个 Add 节点加一个 Mul 节点。但如果推理时这个计算只被用一次,合并成一个 Fused 算子(比如 Scale)就能省掉一次内存读写的开销。这类优化 ONNX Runtime 在加载模型时会自动做一部分,不需要你手动干预。
第二类是算子融合。这是推理优化的大头。常见的融合模式包括 Conv+BN 融合、Conv+ReLU 融合、MatMul+Add 融合等。这些融合能大幅减少访存次数,因为中间结果不用再写回内存再从内存读出来了。ONNX Runtime 和 TensorRT 都在做这件事,只是策略和粒度不一样。
第三类是精度裁剪。把模型里的数据类型从 FP32 改成 FP16 或 INT8。这个改动可以在 ONNX 模型层面做,也可以在推理引擎层面做。ONNX 本身支持不同数据精度的描述,但实际转换时需要特别注意量化工具的选择和数值范围的控制。
4.2 量化:INT8 量化如何加快推理速度
提到 ONNX 模型的优化,量化一定是最受关注的话题之一。量化(Quantization)的本质是用更低精度的数据类型去近似表示原模型的权重和激活值,从而换取更低的计算开销和更小的内存占用。最常见的量化是 INT8 量化,把模型的权重从 FP32 压缩到 INT8,模型大小直接缩小到原来的四分之一,推理延迟通常能降低一半以上,内存带宽的压力也大幅缓解。
ONNX Runtime 的量化方案分两种基础路线:静态量化和动态量化。
动态量化(Dynamic Quantization)是在推理时动态计算激活值的量化范围。它不需要预先准备校准数据集,使用起来最省事,但推理速度提升的幅度相对有限,更适合 CPU 环境下的部署场景。具体做法是用 onnxruntime.quantization 模块里的 quantize_dynamic 接口:
from onnxruntime.quantization import quantize_dynamic, QuantType model_path = "simple_cnn.onnx" quantized_model_path = "simple_cnn_int8.onnx" quantize_dynamic( model_path, quantized_model_path, weight_type=QuantType.QInt8 )静态量化(Static Quantization)需要先用一批有代表性的输入数据(校准数据集)去"观察"激活值的分布范围,然后基于这个范围把激活值也量化成 INT8。这个过程类似一门综合课的期末考试——考试范围(量化参数)是提前划定的,考试时按这个范围去答题(推理)。静态量化的加速效果比动态量化更明显,尤其在 GPU 上,但操作流程更复杂,而且校准数据集的选择直接影响量化后的精度损失。如果校准数据分布和线上真实数据差异太大,量化后的模型精度可能崩得一塌糊涂。
从实战经验来看,INT8 量化有两种情况特别值得做:一种是在 CPU 上部署,目标机器没有强 GPU,量化是提升吞吐最直接的手段;另一种是大模型场景,模型文件动辄好几 GB,量化后能极大缓解显存和内存压力,配合 KV Cache 优化还能进一步扩大并发能力。
量化不是没有代价的。最明显的问题就是精度损失。对分类、检测这类任务,INT8 量化通常能控制在可接受的范围(top-1 准确率下降 1% 以内);但对生成式模型或者对数值敏感的任务(比如车牌识别、OCR 识别),量化后的结果可能出现字符识别错误率升高的问题。我做过的项目里出现过车牌识别模型量化后,字符"5"频繁被识别成"6"的情况。所以量化后的模型一定要做充分的精度回测,而不是只看推理速度。
4.3 ONNX 转 TensorRT:GPU 部署的进阶路线
如果目标部署环境是 NVIDIA GPU,那 ONNX 模型的下一步通常不是直接用 ONNX Runtime,而是转成 TensorRT 引擎。TensorRT 是 NVIDIA 针对自家 GPU 做深度优化的推理引擎,里面有很多黑魔法级的优化(比如 kernel 自动调优、显存复用、低精度推理等)。
ONNX 转 TensorRT 有两种常见方式:一种是用 TensorRT 自带的 trtexec 命令行工具,另一种是在代码里用 TensorRT 的 ONNX Parser 直接解析。比较简单的做法是用 trtexec:
trtexec --onnx=simple_cnn.onnx \ --saveEngine=simple_cnn.engine \ --fp16--fp16 表示开启 FP16 精度推理,这种配置在 Ampere 及以上架构的 GPU 上效果尤其明显。TensorRT 转换时会有一个构建(Build)阶段,这个阶段会分析 ONNX 图结构、自动选择最优的 kernel 实现,还可能做一些层融合的优化。构建过程可能比较慢,几次甚至十几分钟都有可能,但构建好的 engine 文件后续加载推理很快。
有一个值得注意的点是,TensorRT 的 engine 文件和硬件绑定。用 A100 构建的 engine 不能直接拿到 T4 上跑,反之亦然。所以你在训练机上构建的 TensorRT 引擎不能直接部署到生产环境的 GPU 上,必须在目标 GPU 型号相同的环境下重新构建。这也是 ONNX 这类中间格式存在的意义之一——ONNX 可以跨硬件迁移,TensorRT 引擎不行。
YOLO12 ONNX 转 TensorRT 的案例现在社区里讨论度很高,核心流程和我上面写的完全一致:YOLO 模型导出 ONNX,然后 trtexec 转 TensorRT,再用 C++/Python 在 5070 这类显卡上做推理测试。需要注意的点主要是 YOLO 的输出层在导出 ONNX 后可能会包含一些非标准算子,需要提前处理或者使用专用的部署仓库。
5. 推理引擎选型:ONNX Runtime、vLLM、Ollama 怎么选
5.1 一张表看懂主流推理引擎的定位
现在模型推理引擎非常多,光是把名字列出来就能吓到新手。但从使用场景上分,其实可以比较清晰地分成几类。下面用表格来做一个整理:
| 推理引擎 | 适用场景 | 优势 | 注意事项 |
|---|---|---|---|
| ONNX Runtime | 通用模型部署、跨平台 | 生态成熟,支持多端(CPU/GPU/移动端),API 简洁 | 极致性能不如专用引擎 |
| TensorRT | NVIDIA GPU 高性能推理 | 推理速度极快,GPU 利用率高 | engine 和硬件绑定,构建时间长 |
| OpenVINO | Intel CPU/GPU/VPU 推理 | Intel 硬件优化好,边缘部署常用 | 非 Intel 硬件优势不明显 |
| vLLM | 大模型(LLM)在线推理 | PagedAttention 省显存,吞吐高 | 主要服务 LLM,对其他模型支持有限 |
| Ollama | 本地一键部署大模型 | 使用门槛极低,拿来就能跑 | 定制化程度低,适合快速体验 |
| TensorFlow Lite | 移动端/嵌入式设备 | 模型小,端侧优化成熟 | 主要服务 TFLite 模型 |
ONNX Runtime 是目前覆盖面最广的,因为它在 ONNX 模型上做了大量优化,而且支持直接加载 ONNX 模型,不需要二次编译和转换。你在开发阶段用 PyTorch 训练,导出 ONNX 后直接用 ONNX Runtime 跑推理,链路是最短的。
TensorRT 适合对性能有极致要求的场景。如果在生产环境上跑的是固定 GPU 型号,并且模型结构比较稳定(不会频繁改动),值得花时间转成 TensorRT。反过来说,如果模型迭代频繁,每次都要重新构建 TensorRT 引擎,那维护成本会比较高,性价比就低了。
vLLM 和 Ollama 是大模型时代的产物。vLLM 做了大量的显存优化(PagedAttention 机制),能显著提升 LLM 的并发处理能力,适合做在线 API 服务。Ollama 胜在零门槛,本地机器装好之后直接拉模型就能跑,很多人在自己电脑上体验大模型就是这个方案。Ollama 底层其实也集成了多个运行时引擎,模型格式也不止 ONNX 一种,但对用户来说是透明封装好的。
5.2 ONNX 在其中的真实位置
明白了这些引擎的定位之后,ONNX 的位置就很清晰了:它不是一个和 TensorRT、vLLM 直接竞争的推理引擎,而是一个"中间表示层"。你可以把它理解为"通用语言",TensorRT 是"专用高速工具",vLLM 是"大模型专用服务器"。在理想的技术架构里,训练框架负责把模型转成 ONNX 或者其他中间格式(比如 Safetensors、GGUF),推理引擎再从这个中间格式转换成自己最擅长的执行方式。
但有一点必须说明白:ONNX 并不是大模型场景中唯一的中间格式,甚至在很多主流 LLM 部署方案里,ONNX 并不是首选路径。HuggingFace 生态更常用的是 Safetensors + PyTorch 原生推理;本地 CPU 推理方案更常用的是 GGUF(配合 llama.cpp);NVIDIA 的方案则倾向于直接把 PyTorch 模型转成 TensorRT-LLM 适配的格式(通常是按层存权重,而不是一个完整计算图)。
那 ONNX 的优势在哪里?就在于它同时具备"通用性"和"可针对性"。
通用性体现在 ONNX 是一个开放标准,不受单一厂商控制,而且支持范围极广。你写一个模型,导出 ONNX 后,可以在 Windows 上用 ONNX Runtime 跑,在 Linux 上用 TensorRT 跑,在浏览器里用 WebAssembly 跑,在 iPhone 上用 CoreML 导入跑。一套模型四处部署,这个收益在跨端场景里是无法忽略的。
针对性体现在 ONNX 已经在上层工具链上衍生出了一整套生态。比如 HuggingFace 的 Optimum 可以直接导出 ONNX 格式的模型,并且配合 ONNX Runtime 做加速推理,很多 transformer 类模型都有现成的 ONNX 导出教程。这意味着你已经训练好的 BERT、GPT、ViT 这类模型,可能不需要写一行额外的代码,就能获得 ONNX 带来的部署便利性。
所以我的选择建议是:如果项目是多端部署(PC + Web + 移动端),或者后端开发团队主要使用 Java/C# 这类非 Python 语言,那 ONNX 是必需的中间桥梁。如果项目纯粹是 LLM 的云端在线服务,且显卡型号固定、追求极致吞吐,那直接走 vLLM 或者 TensorRT-LLM 路线可能更合适,ONNX 作为兜底方案存在即可。
5.3 多语言部署的典型场景:Java 调用 ONNX 模型
在众多实际需求中,Java 调用 ONNX 模型是非常高频的一个场景。原因在于很多公司的后端服务体系是 Java 写的,模型推理如果依赖 Python 服务,就需要单独部署和管理一个 Python 进程,增加了运维的复杂度。而 ONNX Runtime 官方提供了 Java 绑定,可以直接在 JVM 里面加载 ONNX 模型并执行推理,这样整个推理链路就能完全糅合进 Java 后端服务里。
以车牌识别为例,这个场景的完整链路一般是:先用 PP-OCRv6(PaddleOCR 的第六版)训练或者下载一个文字检测+识别模型,转成 ONNX 格式,然后用 ONNX Runtime Java API 去加载。之所以 PP-OCRv6 能转成 ONNX,是因为 PaddleOCR 的 PaddlePaddle 框架本身支持导出 ONNX,导出的 ONNX 模型可以直接用 ONNX Runtime 的 Java API 跑。
ONNX Runtime 的 Java API 用法也很直接:
import ai.onnxruntime.*; public class OnnxInference { public static void main(String[] args) throws Exception { OrtEnvironment env = OrtEnvironment.getEnvironment(); OrtSession session = env.createSession("simple_cnn.onnx", new OrtSession.SessionOptions()); try (OrtSession.Result result = session.run(env.createTensor(inputData)) ) { // 处理结果 OnnxTensor output = (OnnxTensor) result.get(0); float[][] outputData = (float[][]) output.getValue(); } } }这样做的好处是显而易见的:不需要跨进程通信,不需要额外部署 Python 服务,Java 后端直接和模型推理在同一个 JVM 进程里完成,延迟更低、运维更简单。这也是 ONNX 在工业界落地时最核心的价值之一。很多传统互联网公司的 CV 类业务(OCR、目标检测、图像分类)都通过这种方式把模型推理融入 Java 微服务体系。
6. 大模型时代的 ONNX:从"图像分类时代"到"LLM 时代"
6.1 大模型时代的部署形态变化
大模型(尤其是 LLM)和传统 CNN 模型最大的区别在于推理模式的改变。传统模型是"一次前向传播出结果",输入一张图片、一段短文本,通过几十层网络计算,输出分类结果。而 LLM 是自回归解码——每次只预测一个 token,然后把预测结果拼接到输入里继续预测下一个,循环往复直到生成结束。这种模式带来的结果是:推理时不仅要做矩阵乘法,还要维护大量的历史状态(KV Cache),访存密集程度远高于计算密集程度。
这种形态的变化对部署提出了非常不一样的要求。传统 ONNX 推理里,模型推理一般只需要把输入张量推给引擎即可;而 LLM 推理涉及输入序列长度动态变化、历史 token 的 KV Cache 管理、beam search 或 sampling 策略选择等一系列问题。这也是为什么大模型时代出现了 vLLM、TensorRT-LLM 这些专门做 LLM 推理的引擎——它们不只是执行计算图,更是在做显存管理、调度优化、批量推理等系统层面的工作。
ONNX 在这个阶段的位置,从"推理执行环境"变成了"模型交换格式"。你可以把 ONNX 当成模型在 PyTorch 训练生态与其他推理引擎之间的接口。比如你想把一个 HuggingFace 上的模型部署到 TensorRT 上,中间大概率会用到 ONNX 作为中介格式。有了这个中间层,模型转换的逻辑就能模块化、标准化,而不是每个引擎都重新实现一套模型解析。
6.2 大模型场景下 ONNX 的落地路径
大模型场景下 ONNX 的实际落地路径通常是这样的——先从 HuggingFace 等模型仓库获取预训练权重,用 Optimum 等工具导出为 ONNX 格式,然后分两条路走:一条是直接交给 ONNX Runtime 做 CPU 或 GPU 推理;另一条是再转换成 TensorRT-LLM 等专用格式做高性能 GPU 推理。HuggingFace 的很多 CLIP、ViT、Whisper、BERT 类模型都有现成的 ONNX 权重可以直接下载。
另外,ONNX Runtime 也推出了专门针对 LLM 的优化方案(比如 onnxruntime-genai 扩展),支持 llama、mistral、phi 等常见架构的 ONNX 模型量化推理。这类方案在 CPU 设备上跑 LLM 时表现还不错,但在顶级 GPU 上的性能距离 vLLM 还是有差距。
我在实际使用中的感受是:如果只是本地调试、或者 CPU 上跑一个小规模的 LLM(比如 7B 以下的量化模型),ONNX Runtime 是一个还算靠谱的选择;但如果要大规模服务用户、追求高并发和低延迟,vLLM 和 TensorRT-LLM 还是主要选择。
6.3 对大模型学习者的建议
这几年"动手学大模型"类的教程特别火,上海交大出过一套相关课程,社区里也有各种学习路线。但说实话,很多所谓的"大模型学习路线"内容安排得太功利,一上来就让人去跑 vLLM、Ollama,却忽略了最基础的知识储备。其中 ONNX 这条链路往往是被轻视的一环,但它恰恰是很多工程师从"会用大模型"跨到"能上手部署大模型"的分水岭。
我给大模型学习者的建议是:不要一上来就陷入"哪个框架最强"的比拼里。先把模型推理的基本链路走通——训练一个模型(哪怕是 MNIST 分类),导出 ONNX,用 ONNX Runtime 跑推理,再用 TensorRT 试一次 GPU 加速。这个过程能帮你建立对"模型从训练到部署"全链路的直觉,比单纯看教程有用得多。
记住一个学习原则:模型格式只是载体,计算图和算子的本质才是根本。你把 ONNX 的图结构搞明白了,后面学 vLLM 的 PagedAttention、TensorRT-LLM 的层融合、Ollama 的量化格式,都会觉得"不过是把同一件事做成了不同的性能版本",底层逻辑是通的。
7. 常见 ONNX 问题和排查思路实录
做 ONNX 相关的部署做久了,遇到的坑基本能总结成一张清单。我这里挑几个出现频率最高的问题,以及对应的排查思路,分享出来供大家参考。
7.1 导出时报"Unsupported operator"错误
这是最常见的问题之一。PyTorch 里的某些算子(尤其是比较新的、不常见的算子)在 ONNX 的算子集合里还没有对应的映射,导出时就会报错。
排查思路:先看错误信息里提示的是哪个算子在搞事。如果是比较新的算子,可以试着升级 PyTorch 和 ONNX 的版本,新版通常会增加更多算子映射(比如 FlashAttention 相关的融合算子)。如果升级后还是不行,就需要考虑修改模型结构来绕过这个算子,比如把自定义算子拆解成多个 ONNX 支持的基础算子。还有一种办法是写自定义算子注册(ONNX 的自定义 op 机制),但这个方案复杂度较高,做之前要评估清楚收益。
7.2 导出成功但推理结果不对
这种情况通常不是 ONNX 本身的问题,而是模型在导出时丢掉了一些信息。最常见的原因是 batch norm 层在训练模式和推理模式下的行为不一样,导出时如果没有调用 model.eval(),batch norm 还是在用训练时的统计方式(用当前 batch 的均值方差),导出后的模型和实际推理行为对不上。
另外一个常见原因是动态 shape 处理不当。如果你的模型在 forward 里有依赖输入 shape 的条件逻辑,导出时用的 dummy input 一旦和线上输入尺寸不一致,导出的计算图就可能走错分支。
排查这类问题时,强烈建议用 Netron 打开导出的 ONNX 模型,人眼检查一遍网络结构。很多时候问题在 Netron 里一眼就能发现,比盯着代码猜要高效得多。
7.3 量化后精度暴跌
INT8 量化后精度下降是正常现象,但"暴跌"就说明哪里出了问题。检查顺序如下:第一,确认校准数据集的规模和质量,一般建议至少几百到几千张有代表性的图片或文本样本,覆盖尽可能多的分布情况;第二,检查有没有不适合量化的层,比如最后的全连接层、带有特殊数学运算的层,必要时可以在量化时对特定的层做精度保护(不量化该层);第三,确认模型的动态范围是否稳定,如果模型输入数据的数值尺度变化非常大(比如 0 到 1 的图片和 0 到 255 的图片混在一起),量化效果肯定会受影响,建议先做输入归一化再量化。
7.4 遇到 Unknown model 或 Unsupported DataType
这类错误通常说明推理引擎的版本和模型使用的 opset 版本不匹配。比如你拿 opset 17 导出的模型,喂给一个只支持 opset 15 的 ONNX Runtime 旧版本,就会报类似错误。解决方式是降低导出的 opset 版本,或者升级推理引擎的版本。
如果是 DataType 不支持,很可能是在模型里引入了自定义的、非标准的数据类型(比如某些特殊场景下的复数运算),这种情况下需要考虑调整模型实现,避免使用推理引擎不支持的中间数据类型。
7.5 我在实操中的几条经验
最后分享几条吃过亏才记住的经验。
第一,导出 ONNX 之前,先在 PyTorch 侧把模型推理的结果保存一份,导出的 ONNX 推理结果和它做精确对比。不要省这一步,它能帮你提前发现百分之八十的模型迁移问题。
第二,保护好自己的 CPU 环境。ONNX Runtime 在 CPU 上的性能优化已经做得比较好了,很多场景(尤其是中小模型)CPU 推理速度和 GPU 差距并没有想象中那么大。急着上 GPU 之前,先试试 CPU 推理的 benchmark,可能会帮你省下一笔显卡预算。
第三,多留意大模型和边缘设备领域的格式演进速度。ONNX 仍然是目前覆盖面最广的中间表示,但 GGUF、Safetensors 等格式在 LLM 领域已经有很强的生态了。做技术选型时,不要有"一个格式走天下"的执念,根据场景选择最合适的链路才是务实的工作方式。
第四,ONNX 模型并不是越大越准,也不是越小越快。过量化会导致精度受损,过度融合也可能引入风险,合理评估精度和性能的平衡点,才是工程上最应该花时间的部分。