PyTorch09-PyTorch模型部署
九、PyTorch模型部署
到目前为止,我们已经初步掌握了如何用 PyTorch 框架实现与训练神经网络模型,也学习如何在训练环境下做推理、将模型文件保存下来。但很多时候,我们在软硬件环境A下训练与保存得到了模型文件,实际使用时可能无法还原当时训练环境A中的软硬件,甚至要求使用与训练环境A完全不同的另一套软硬件环境。此时我们可能会想,有没有一种类似于 Java 虚拟机一样的东西能让编译好的 Java 字节码到处运行,我们只要在实际使用中将恢复或搭建一个兼容的 Java 虚拟机 环境即可运行已有的 Java 字节码(其实就是 Java 程序的一种保存形式,是其编译得到的内容)。
可喜地是,在神经网络模型的生态环境内,确实存在一个类似于 Java
虚拟机环境的开源解决方案,它就是 ONNX
,它将使用PyTorch框架实现的神经网络模型 转换成 .onnx
格式文件。与此同时,与之配套使用的还需要解析与运行 .onnx
格式文件的推理引擎及加速硬件。
本节学习的主要内容包含:ONNX的介绍与使用、推理引擎ONNX Runtime的介绍与使用、其他推理引擎。
9.1 ONNX的介绍与使用
9.1.1 ONNX的介绍
ONNX 官网:https://onnx.ai/ 。 ONNX
GitHub代码仓库地址:https://github.com/onnx/onnx 。
ONNX(Open Neural Network Exchange,开放式神经网络交换)
是由微软、Meta
等机构在2017年共同创立的开源AI生态标准,旨在建立一个通用的模型文件格式。它充当了不同深度学习框架(如
PyTorch、PaddlePaddle
等)与各类硬件推理引擎之间的“跨平台桥梁”。通过将训练好的神经网络模型导出为统一的
.onnx
计算图格式,开发者能够实现“一次训练,到处部署”,轻松将模型运行在 ONNX
Runtime、TensorRT、华为昇腾Asend 等多样化的软硬件生产环境中。
目前支持将神经网络模型转换成 ONNX 格式的深度学习框架,如下图所示:

ONNX 本质上只是一个“中间格式”(IR,Intermediate Representation),它就像是一份用统一标准绘制的建筑蓝图。这份蓝图本身是不能直接在硬件上运行的,它必须通过各家硬件厂商自己开发的推理引擎或编译器来进行解析、编译和执行。
ONNX 目前得到了 广泛支持,并被应用于众多框架、工具和硬件中。它能够实现不同框架之间的互操作性,并简化从研究到生产的流程,从而加速人工智能领域的创新,支持ONNX的推理引擎如下图所示:

基于ONNX标准,硬件厂商与软件厂商(一般来说提供加速硬件设施的硬件厂商也会同时开发出对应的推理引擎)只需要在自己的硬软件环境中实现 ONNX 文件的解析、运行与性能优化,就能把用其他深度学习框架开发的神经网络模型运行在自己自家软硬件产品上,实现生态融合与产品推广。
9.1.2 ONNX 的安装
ONNX 发布的软件包已发布在 PyPi 上,其官方推荐通过 pip 安装ONNX:
1
pip install onnx
详细的安装说明,包括常用构建选项和常见错误,请点击此处查看。
9.1.3 ONNX 的使用
9.1.3.1 模型转换为ONNX格式
使用传统的 torch.onnx.export 导出模型时的示例代码如下:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23import torch
model = MyModel()
# 导出模型前,必须调用model.eval()或者model.train(False)
# 避免因为`dropout`或`batchnorm`等运算符在推理和训练模式下的不同产生错误
model = model.eval()
# 使用传统的 `torch.onnx.export` 导出模型时,必须准备一个符合输入尺寸和类型的示例输入
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(
model,
dummy_input,
"model.onnx",
export_params=True,
opset_version=13, # ONNX 算子集版本,13 = 2020
input_names=["input"], # 输入节点名
output_names=["output"], # 输出节点名
dynamic_axes={ # 支持动态维度
"input": {0: "batch_size", 2: "height", 3: "width"},
"output": {0: "batch_size"}
}
)
9.1.3.2 ONNX模型的检验
”.onnx“格式文件导出后,可以通过如下方法校验其是否可用:
1 | |
9.1.3.3 ONNX可视化
9.1.3.3.1 Netron 介绍
前面所述,通过TensorBoard可以对模型结构与输出特征图进行可视化。onnx将深度学习模型导出为".onnx"格式文件后,也有一个类似的开源项目 Netron 可以对".onnx"格式代表的模型进行可视化。
Netron 是一个用于查看神经网络、深度学习和机器学习模型的工具。
Netron 支持 ONNX、TensorFlow Lite、PyTorch、torch.export、ExecuTorch、TorchScript、TensorFlow、Core ML、OpenVINO、Keras、Caffe、Darknet、Safetensors 和 NumPy。Netron 对 MLIR、JAX、GGUF、RKNN、ncnn、MNN、PaddlePaddle 和 scikit-learn 提供实验性支持。
9.1.3.3.2 Netron 安装
Netron 不仅支持通过pip安装,且在不同操作系统下提供了对应的安装文件。如下:
Python: pip install netron, then run
netron [FILE] or netron.start('[FILE]').
macOS: Download
the .dmg file or run
brew install --cask netron.
Linux: Download
the .deb or .rpm file.
Windows: Download
the .exe installer or run
winget install -s winget netron.
甚至,Netron 官网还提供了一个web界面,可以直接可视化".onnx"格式代表的模型。
9.1.3.3.2 Netron 使用
可用浏览器直接访问https://netron.app/打开netron官网提供的在线可视化界面,打开本地甚至线上的".onnx"文件,即可得到可视化效果。
下图来自Netron Github代码仓库中提供示例的实际效果:Netron可视化squeezenet1.0-3.onnx:

对于自定义的模型导出为".onnx"文件保存到本地后,可以通过如下方法可视化:
1
2
3
4
5
6
7
8
9
10
11
12
13# 10-使用netron可视化onnx格式模型.py
from pathlib import Path
import netron
# 假设 ./squeezenet1.0-3.onnx 就是通过ONNX转换得到的 .onnx 文件路径。
ONNX_FILE = "./squeezenet1.0-3.onnx"
# Netron 内置 HTTP 服务绑定的本机地址和端口;127.0.0.1 仅允许本机浏览器访问。
ADDRESS = ("127.0.0.1", 8081)
# Netron 把服务放在“守护线程”中;若 start() 后脚本立即结束,服务线程也会
netron.start(file=str(ONNX_FILE), address=ADDRESS, browse=True)
# 持续等待,让主线程不退出,从而让 Netron 的 HTTP 服务持续监听端口。
netron.wait()
再细致查看可以发现,Netron 提供的模型可视化界面中,不仅可以对每个”节点“单独查看,甚至能查看”节点“中相关属性值如卷积核的形状、边缘、步长,以及”节点“输入如模型自身的权重 W 和 偏置 B ,还有输出。通过这个可视化结果,可以非常方便直观地查看模型的结构以及数据的变化过程。
此外可视化结果界面左侧还提供了一些”调节性“功能按钮如”Show Attribute“、”Show Weights“,还可以把即时呈现的可视化效果以PNG/SVG格式持久化保存。
9.2 与onnx配套的常用推理引擎环境
9.2.1 模型训练到跨平台部署
下图展示了”深度学习框架开发神经网络模型(训练阶段)、将各框架的模型统一导出为中间表示(桥梁与标准化阶段)、具体的推理引擎解析运行模型的中间表示(生产与运行阶段)“
这三个阶段的示意与经典代表,描绘了现代深度学习模型从训练到跨平台部署的完整生命周期与流水线(Pipeline):

上图的具体解释如下: 1. 深度学习框架(左侧:训练阶段)
- 代表工具:PyTorch、TensorFlow、Caffe、MxNet 等。
- 核心任务:在这个阶段,开发者使用各种主流框架编写代码、设计网络结构,并在服务器或集群上使用海量数据进行模型训练(Training),获取具有预测能力的权重。
- 中间表示 / IR(中间:桥梁与标准化阶段)
- 代表格式:ONNX、TorchScript、Caffe 内部格式等。
- 核心任务:由于不同的训练框架格式不通,无法直接在生产环境中高效运行,因此需要通过“转换”将各框架的模型统一导出为中间表示(Intermediate Representation, IR)。其中,ONNX 就是最核心、最通用的标准中间格式。在这个阶段还可以对计算图进行初步的优化。
- 推理引擎(右侧:生产与运行阶段)
- 代表工具:TensorRT、ONNX Runtime、PaddleInference、NCNN、OpenVINO 等。
- 核心任务:将标准化的中间表示(如 ONNX 文件)输入到具体的推理引擎中。推理引擎会针对不同的硬件(CPU、GPU、NPU 等)进行深度的算子融合与编译优化,最终在生产环境中高效地运行(Inference)推理任务。
9.2.2
ONNX Runtime的介绍与使用
9.2.2.1 ONNX Runtime的介绍
ONNX 官网:https://onnxruntime.ai/ 。 ONNX
GitHub代码仓库地址:https://github.com/microsoft/onnxruntime
。
ONNX Runtime是一个跨平台的推理和训练机器学习的加速器。
ONNX Runtime 推理能够加快用户体验并降低成本,它支持来自
PyTorch 和 TensorFlow/Keras 等深度学习框架以及
scikit-learn、LightGBM、XGBoost 等传统机器学习库的模型。ONNX Runtime
与不同的硬件、驱动程序和操作系统兼容,并通过在适用情况下利用硬件加速器以及图优化和转换来提供最佳性能。
ONNX Runtime一般与ONNX搭配使用,将由ONNX导出的".onnx"格式的文件输入给ONNX Runtime这一推理引擎,实现不同深度学习框架模型的到处部署与推理。
9.2.2.2 ONNX Runtime的安装
官方提供的安装文档:https://onnxruntime.ai/docs/install/。
ONNX 与 ONNX Runtime 的兼容性矩阵见: ONNX 与 ONNX Runtime 的兼容性矩阵。
1 | |
其他硬件加速设备、不同开发语言和应用场景等下的相关安装细节,请参考https://onnxruntime.ai/docs/install/。
ONNX 与 ONNX Runtime 兼容核心规则
- 向后兼容(Backwards Compatibility):所有版本的 ONNX Runtime (ORT) 都能运行早期标准下导出的模型(支持低至 opset 7 及以上的模型)。 1, 2
- 向前兼容/上限约束(Forward Compatibility):模型的 opset 版本 不能超过当前使用的 ONNX Runtime 版本所支持的最高 opset 范围。如果导出的 opset 高于运行时支持的上限,会直接抛出模型加载失败错误。 1, 2
- 独立演进机制:ONNX 规范(Opset / IR 版本)与 ONNX
Runtime 软件版本按不同周期独立迭代(ONNX 通常约4个月发布一次,ORT
按季度发布)。 1
#### 9.2.2.3
ONNX Runtime的使用
9.2.2.3.1 使用流程讲解
ONNX Runtime的使用流程大致如下:
- 根据转换得到的.onnx文件创建一个ONNX Runtime 推理器;
- 构建onnx的输入,要注意输入的名字与数据类型;
- 执行模型推理。
以下示例代码,展示了使用ONNX Runtime运行转换得到的".onnx"模型文件,并进行推理的核心过程:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23# 导入onnxruntime
import onnxruntime
# 需要进行推理的onnx模型文件名称
onnx_file_name = "model.onnx"
# onnxruntime.InferenceSession用于获取一个 ONNX Runtime 推理器
ort_session = onnxruntime.InferenceSession(onnx_file_name)
# 构建字典的输入数据,字典的key需要与我们构建onnx模型时的input_names相同
# 输入的input_img 也需要改变为ndarray格式
input_img = rng.standard_normal(size=(1, 3, 224, 224)).astype(np.float32)
ort_inputs = {'input': input_img}
# 我们更建议使用下面这种方法,因为避免了手动输入key
# ort_inputs = {ort_session.get_inputs()[0].name:input_img}
# run是进行模型的推理
# 第一个参数是“希望取回哪些输出”。None 表示取回模型的全部输出;
# 第二个参数为构建的输入值的字典:{输入名称: NumPy 数组}
# 由于返回的结果被列表嵌套,因此我们需要进行[0]的索引
ort_output = ort_session.run(None, ort_inputs)[0]
# output = {ort_session.get_outputs()[0].name}
# ort_output = ort_session.run([output], ort_inputs)[0]
在上述示例代码中,有几个需要注意的点:
ONNX Runtime公支持opset 7及以上的 .onnx 模型。如果使用官方提供的 squeezenet1.0-3.onnx 示例文件 直接替换上述示例代码中上的 onnx_file_name ,运行时会报错,提示NOT_IMPLEMENTED : Could not find an implementation for Relu(1) node with name ''。此时可做个转换,将原squeezenet1.0-3.onnx转换并保存为支持opset 7及以上的 .onnx 模型文件squeezenet1.0-3_opset7.onnx。 转换过程有需要再研究。- PyTorch模型的输入为tensor,而ONNX的输入为array。构造输入时要注意数据类型。
- 输入ONNX的array的shape应该和我们导出模型的输入的shape相同。
- 在构建输入的字典时,我们需要注意字典的key应与导出ONNX格式设置的input_name相同,因此我们更建议使用上述的第二种方法构建输入的字典,即自动获取导出模型时用到的输入名。
9.2.2.3.2 完整可运行示例
以下示例整理自 PyTorch 官方教程:export_simple_model_to_onnx_tutorial
此完整示例代码可正常执行,其实现的主要功能与流程如下:
- 定义一个简单的图像分类模型;
- 通过 torch.onnx.export() 导出为 ONNX;
- 用 onnx.checker 检查导出的模型;
- 用 ONNX Runtime 执行同一份输入;
- 比较 PyTorch 与 ONNX Runtime 的数值输出。
- 执行时,会在当前目录生成一个""文件,此外命令行输出内容:
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190# 13-自定义模型-用ONNX导出-ONNXRuntime运行-完整示例.py
"""自定义 PyTorch 图像分类模型 -> ONNX -> ONNX Runtime 的完整最小示例。
运行方式(在项目根目录执行):
python ./13-自定义模型-用ONNX导出-ONNXRuntime运行-完整示例.py
所需依赖:
python -m pip install onnx onnxruntime
官方当前推荐的新导出器还需要 onnxscript:
python -m pip install onnxscript
本文件会优先尝试教程中的 ``dynamo=True`` 新导出器;若缺少 onnxscript,则自动回退到 ``dynamo=False`` 的传统
导出器,以保证“导出 -> 推理 -> 对比”的学习流程依然可以完整运行。安装 onnxscript 后,脚本会自动优先采用推荐的新导出器。
"""
# importlib.util 用于判断 onnxscript 是否已经安装,不会真正导入该包。
import importlib.util
# Path 用于构造与脚本同目录的 ONNX 输出路径,不依赖终端当前工作目录。
from pathlib import Path
# torch.nn 提供层类;torch.nn.functional 提供函数形式的 ReLU、最大池化等操作。
import torch
import torch.nn as nn
import torch.nn.functional as F
# onnx 用来读取导出后的文件,并检查其图结构是否有效。
import onnx
# onnxruntime 用来脱离 PyTorch 执行 ONNX 图。
import onnxruntime as ort
# 将导出的模型放到本脚本的同一目录。每次运行会覆盖此前同名的演示输出。
ONNX_FILE = Path(__file__).resolve().with_name("13-image_classifier_model.onnx")
class ImageClassifierModel(nn.Module):
"""与官方教程相同的、输入为单通道 32×32 图像的 10 类分类网络。
其结构近似 LeNet:两组“卷积 + ReLU + 最大池化”提取局部特征,再用三层全连接层
输出 10 个类别的原始分数(logits)。模型随机初始化,仅用来演示导出流程,尚未训练,
因而其类别结果没有实际语义。
"""
def __init__(self) -> None:
"""创建模型中可学习的卷积层和全连接层。"""
# 初始化 nn.Module 内部管理参数、子模块等所需的数据结构。
super().__init__()
# 输入是 1 通道灰度图;6 个 5×5 卷积核会产生 6 张特征图。
self.conv1 = nn.Conv2d(in_channels=1, out_channels=6, kernel_size=5)
# 第二层将 6 通道特征图变换成 16 通道特征图。
self.conv2 = nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5)
# 32×32 输入经过两次 5×5 卷积和 2×2 池化后,空间尺寸为 5×5,
# 故展平后特征数为 16 * 5 * 5 = 400。
self.fc1 = nn.Linear(in_features=16 * 5 * 5, out_features=120)
# 将 120 维特征继续映射为 84 维。
self.fc2 = nn.Linear(in_features=120, out_features=84)
# 最后一层输出 10 个 logits;这里不使用 softmax,分类训练时通常会将 logits
# 直接交给 CrossEntropyLoss,推理时若需要概率再单独做 softmax。
self.fc3 = nn.Linear(in_features=84, out_features=10)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""定义输入 Tensor 在各层之间的前向计算过程。"""
# 输入形状:(N, 1, 32, 32) -> 卷积后 (N, 6, 28, 28) -> 池化后 (N, 6, 14, 14)。
x = F.max_pool2d(F.relu(self.conv1(x)), kernel_size=(2, 2))
# 卷积后 (N, 16, 10, 10) -> 池化后 (N, 16, 5, 5)。
x = F.max_pool2d(F.relu(self.conv2(x)), kernel_size=2)
# 从 dim=1 开始展平,保留批次维 N:(N, 16, 5, 5) -> (N, 400)。
x = torch.flatten(x, start_dim=1)
# 前两层全连接层后都使用 ReLU 引入非线性。
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
# 输出形状为 (N, 10),每一行是一个样本的 10 个类别 logits。
return self.fc3(x)
def export_model_to_onnx(
torch_model: nn.Module,
example_inputs: tuple[torch.Tensor, ...],
output_path: Path,
) -> str:
"""优先使用官方推荐的新导出器;缺少其依赖时使用传统导出器。
返回值是实际采用的导出方式,仅用于在终端中明确展示本次运行走的是哪条路径。
"""
# PyTorch 2.5 起推荐 dynamo=True。它基于 torch.export 与 FX 捕获模型图,
# 但还依赖未随 torch 自动安装的 onnxscript。
has_onnxscript = importlib.util.find_spec("onnxscript") is not None
if has_onnxscript:
# 这对应官方 HTML 教程的核心代码:先得到 ONNXProgram,再调用 save() 写文件。
onnx_program = torch.onnx.export(
torch_model,
example_inputs,
dynamo=True,
input_names=["image"],
output_names=["logits"],
)
onnx_program.save(output_path)
return "推荐的新导出器:dynamo=True"
# 仅在缺少 onnxscript 时进入此分支。dynamo=False 使用传统 TorchScript 导出器;
# 它目前仍可用,但 PyTorch 新项目应在安装 onnxscript 后优先使用上面的 dynamo=True。
print("提示:未检测到 onnxscript,改用 dynamo=False 的传统导出器。")
print(" 安装 onnxscript 后,下次运行会自动使用官方推荐的 dynamo=True。")
torch.onnx.export(
torch_model,
example_inputs,
f=str(output_path),
dynamo=False,
input_names=["image"],
output_names=["logits"],
# 使用较新的 opset 18,当前 ONNX Runtime 可稳定运行。
opset_version=18,
)
return "传统导出器:dynamo=False(因为未安装 onnxscript)"
def main() -> None:
"""执行“自定义模型 -> ONNX -> ONNX Runtime -> 数值对比”的全部步骤。"""
# 固定随机种子,让模型随机权重、示例输入及最终打印结果在每次运行时可复现。
torch.manual_seed(0)
# 实例化教程中的自定义模型。eval() 切换到推理模式;本模型没有 Dropout/BatchNorm,
# 但实际导出模型时仍应养成在导出前调用 eval() 的习惯。
torch_model = ImageClassifierModel().eval()
# 导出器需要一份“示例输入”来记录 forward 中的张量形状与运算路径。
# 即使模型只有一个输入,args 也必须是元组,因此末尾的逗号不可省略。
example_inputs = (torch.randn(1, 1, 32, 32),)
# 导出为 ONNX 二进制 protobuf 文件,并记录实际使用的导出器。
exporter_name = export_model_to_onnx(torch_model, example_inputs, ONNX_FILE)
# 再将刚保存的 ONNX 文件读回内存;checker 会检查图、节点、类型等结构是否合法。
onnx_model = onnx.load(ONNX_FILE)
onnx.checker.check_model(onnx_model)
# ONNX Runtime 接口使用 NumPy 数组,而非 PyTorch Tensor。detach() 使 Tensor 脱离
# 自动求导图,cpu() 确保数据在 CPU,numpy() 才能返回 ndarray。
onnx_inputs = [tensor.detach().cpu().numpy() for tensor in example_inputs]
# 建立 ONNX Runtime 推理会话。CPUExecutionProvider 表示无论是否有 GPU,都使用 CPU。
ort_session = ort.InferenceSession(
str(ONNX_FILE),
providers=["CPUExecutionProvider"],
)
# get_inputs() 返回 ONNX 图期望的输入名称、形状和类型。用 zip() 动态组装字典,
# 不手写 "image",便能避免模型输入名称变化时字典键名写错。
onnxruntime_input = {
input_arg.name: input_value
for input_arg, input_value in zip(ort_session.get_inputs(), onnx_inputs)
}
# print(f"onnxruntime_input: {onnxruntime_input}")
# run(None, ...) 的 None 表示取回模型全部输出。返回值是 list;本模型只有一个输出,
# 故通过 [0] 取出形状为 (1, 10) 的 NumPy 数组。
onnxruntime_output = ort_session.run(None, onnxruntime_input)[0]
# 使用完全相同的输入调用原 PyTorch 模型;inference_mode() 不建立反向传播所需计算图。
with torch.inference_mode(): # 或 torch.no_grad():
torch_output = torch_model(*example_inputs)
# 将 ONNX Runtime 的 ndarray 转为 Tensor,再逐元素比较。assert_close 默认会同时检查
# 形状、数据类型和数值误差;若导出或推理发生错误,此处会抛出清晰的 AssertionError。
onnxruntime_output_tensor = torch.from_numpy(onnxruntime_output)
torch.testing.assert_close(torch_output, onnxruntime_output_tensor)
# 输出关键元信息和对比结果。注意模型未训练,10 个 logits 只用于验证数值一致性,
# 不能把最大的 logit 当成有实际意义的分类预测。
print(f"PyTorch 版本:{torch.__version__}")
print(f"ONNX 版本:{onnx.__version__}")
print(f"ONNX Runtime 版本:{ort.__version__}")
print(f"ONNX 导出方式:{exporter_name}")
print(f"已导出 ONNX 文件:{ONNX_FILE}")
print(f"ONNX 模型输入:{[(item.name, item.shape, item.type) for item in ort_session.get_inputs()]}")
print(f"ONNX 模型输出:{[(item.name, item.shape, item.type) for item in ort_session.get_outputs()]}")
print(f"PyTorch 输出形状:{tuple(torch_output.shape)}")
print(f"ONNX Runtime 输出形状:{onnxruntime_output.shape}")
print(f"最大绝对误差:{(torch_output - onnxruntime_output_tensor).abs().max().item():.8f}")
print("PyTorch 与 ONNX Runtime 的输出一致!")
print(f"示例输出 logits:{onnxruntime_output}")
# print(f"示例输出 logits 的最大概率对应分类索引: {onnxruntime_output.argmax()}")
if __name__ == "__main__":
main()1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21[torch.onnx] Obtain model graph for `ImageClassifierModel([...]` with `torch.export.export(..., strict=False)`...
[torch.onnx] Obtain model graph for `ImageClassifierModel([...]` with `torch.export.export(..., strict=False)`... ✅
[torch.onnx] Run decompositions...
[torch.onnx] Run decompositions... ✅
[torch.onnx] Translate the graph into ONNX...
[torch.onnx] Translate the graph into ONNX... ✅
[torch.onnx] Optimize the ONNX graph...
[torch.onnx] Optimize the ONNX graph... ✅
PyTorch 版本:2.14.1
ONNX 版本:1.23.2
ONNX Runtime 版本:1.31.0
ONNX 导出方式:推荐的新导出器:dynamo=True
已导出 ONNX 文件:/Users/jsy/Dev/projects/python/2026年7月20日开始学习PyTorch/source/09第九章/13-image_classifier_model.onnx
ONNX 模型输入:[('image', [1, 1, 32, 32], 'tensor(float)')]
ONNX 模型输出:[('logits', [1, 10], 'tensor(float)')]
PyTorch 输出形状:(1, 10)
ONNX Runtime 输出形状:(1, 10)
最大绝对误差:0.00000004
PyTorch 与 ONNX Runtime 的输出一致!
示例输出 logits:[[-0.10066912 -0.0901743 0.03644537 0.05913131 -0.04214709 0.02031167
0.03142105 0.07241504 -0.02758684 0.12333921]]
9.2.2.3.3 参考文档
- 使用ONNX进行部署并推理
- ONNX模型使用与多环境部署
- 10-4-模型部署实战.md
- PyTorch Tutorials,此为PyTorch官方教程,其中提供了很多示例。
9.2.3 其他推理引擎
前面介绍的 ONNX Runtime 是运行 ONNX
模型的一种通用推理引擎,但并不是唯一选择。实际部署时还会遇到其他推理运行时、模型格式和编译工具链。
这三类概念很容易混淆,先作一个区分:
| 类型 | 作用 | 例子 |
|---|---|---|
| 推理运行时 / 引擎 | 读取模型、调度算子并在目标硬件上执行推理 | ONNX
Runtime、TensorRT、OpenVINO、ExecuTorch、llama.cpp |
| 模型格式 / 中间表示 | 保存计算图、权重或二者的信息,供某个运行时读取 | ONNX、Core ML、.tflite、GGUF、OpenVINO IR |
| 编译器或编译基础设施 | 将模型图进一步优化、编译为特定硬件可执行的程序或库 | torch.compile/Inductor、AOTInductor、IREE、Apache
TVM |
因此,ONNX 并不等同于
ONNX Runtime:前者是模型的中间表示,后者是能够执行该表示的一种运行时。对于常见的计算机视觉、语音和传统机器学习模型,ONNX
仍是非常成熟的跨框架交换格式;但它不是所有场景下的唯一答案,尤其是端侧部署和大语言模型部署往往需要结合目标硬件与运行时重新选择。
9.2.3.1 PyTorch 原生生态路线
当训练、部署都以 PyTorch 为主,并且不急于把模型交给其他框架或语言时,可以优先考虑 PyTorch 原生方案。
torch.compile与 Inductor- 是什么:
torch.compile()是 PyTorch 2.x 提供的图捕获与编译入口,默认通常使用 Inductor 生成并优化后端代码。 - 适合什么:希望在现有 Python/PyTorch 推理代码上较少改动地获得优化时,可以先尝试它。
- 要注意什么:它主要优化“当前 PyTorch 进程中的执行”,不是把模型保存为一个可被任意运行时读取的通用模型文件;模型可导出性、首次编译开销和实际加速效果都应在目标环境中测试。
- 是什么:
AOTInductor(Ahead-Of-Time Inductor)
- 是什么:AOTInductor 可将通过
torch.export得到的计算图预先编译并打包,面向不希望在部署端再进行即时编译的场景。 - 适合什么:部署环境与训练框架仍围绕 PyTorch,但希望使用预编译产物、减少部署端 Python 依赖或启动期编译开销时。
- 要注意什么:它更接近 PyTorch 自身的部署工具链,而不是 ONNX 的通用替代品;具体支持的硬件、Python/C++ 调用方式应以对应 PyTorch 版本文档为准。
- 是什么:AOTInductor 可将通过
ExecuTorch
- 是什么:ExecuTorch 是 PyTorch
面向端侧(edge)的运行时与部署工具链,可将模型经
torch.export等流程处理后交给移动端、嵌入式设备等受资源约束的平台执行。 - 适合什么:目标是 Android、iOS、嵌入式 Linux 或特定边缘硬件,并且希望保持 PyTorch 模型开发流程时。
- 要注意什么:它的目标是端侧运行,模型算子覆盖、量化方案和硬件后端与服务器端 PyTorch/ONNX Runtime 不完全相同,部署前需要检查目标后端是否支持模型中的算子。
- 是什么:ExecuTorch 是 PyTorch
面向端侧(edge)的运行时与部署工具链,可将模型经
9.2.3.2 大语言模型(LLM)与生成式 AI
LLM 推理除了模型前向计算外,还需要处理分词器、连续生成、批处理调度和 KV Cache(保存已生成 token 的注意力键值)等问题。因此,除了通用 ONNX 路线外,也出现了针对 LLM 的格式与运行时。
GGUF 与
llama.cpp- GGUF 是什么:GGUF 是常见于
llama.cpp生态的模型文件格式,通常将量化后的权重、模型元数据和分词器相关信息放在一个文件中。 llama.cpp是什么:llama.cpp才是读取 GGUF 并执行推理的运行时/项目;GGUF 本身不是推理引擎,也不是像 ONNX 那样的通用计算图交换标准。- 适合什么:在个人电脑、CPU 或消费级 GPU
上离线运行经过量化的 LLM,且更重视内存占用和易用性时,GGUF +
llama.cpp是常见选择。
- GGUF 是什么:GGUF 是常见于
Safetensors 与 Transformers 生态
- Safetensors 是什么:Safetensors
是安全、高效地保存张量权重的文件格式,常与 Hugging Face
transformers一起使用。 - 它不是什么:它通常不保存完整计算图,也不提供推理执行能力,因此不能单独视为 ONNX、GGUF 或 ONNX Runtime 的替代品。
- 如何运行:还需要与模型配置、分词器和实际运行模型的代码/运行时配合,例如
transformers、vLLM、SGLang 或 TensorRT-LLM 等。
- Safetensors 是什么:Safetensors
是安全、高效地保存张量权重的文件格式,常与 Hugging Face
LLM 服务端运行时
- vLLM、SGLang:面向服务端高并发生成场景,重点处理连续批处理、KV Cache 管理与吞吐量;它们不是模型文件格式。
- TensorRT-LLM:面向 NVIDIA GPU 的 LLM 优化工具链与运行时,适合对 NVIDIA 平台性能有较高要求的部署场景。
- ONNX 仍可使用:ONNX 并没有因为 LLM 而失效;例如 ONNX Runtime 也提供面向生成式 AI 的相关能力。实际选择应由模型规模、目标硬件、量化方案和服务并发需求共同决定。
9.2.3.3 面向特定平台的运行时与模型格式
这类方案通常能更充分地利用某家厂商的硬件,但可移植性一般弱于 ONNX。
| 方案 | 主要目标平台 | 适合场景与说明 |
|---|---|---|
Core ML(.mlmodel /
.mlpackage) |
Apple 的 iOS、iPadOS、macOS、watchOS 等 | Apple 平台的官方机器学习框架。通过 Core ML Tools 等工具转换模型后,可由系统根据设备能力调度 CPU、GPU 与 Apple Neural Engine。 |
LiteRT / TensorFlow
Lite(.tflite) |
Android、嵌入式及移动端 | LiteRT 是 TensorFlow Lite 演进后的名称;.tflite
仍是常见模型文件扩展名。它主要服务 TensorFlow/Keras
生态,也常用于移动端量化模型。 |
OpenVINO 与 OpenVINO IR(.xml /
.bin) |
Intel CPU、GPU、NPU 与部分 VPU | Intel 的推理工具链。OpenVINO 可以读取 ONNX 等模型;在其部署流程中也常使用 IR 形式与设备相关优化。 |
| TensorRT | NVIDIA GPU | NVIDIA 的高性能推理优化器与运行时,可从 ONNX 等来源构建 TensorRT engine;适合模型与部署硬件均固定为 NVIDIA GPU 的场景。 |
平台专用方案并非“绝对更快”。能否加速以及加速幅度会受到模型结构、输入大小、量化方式、设备型号和运行时版本影响,应使用真实模型在目标设备上基准测试。
9.2.3.4 编译器级开源方案
这类项目更接近“把模型编译成目标设备程序”的工具链。它们通常比直接使用 ONNX Runtime 更灵活,但接入、调试和硬件适配成本也更高。
MLIR
- MLIR(Multi-Level Intermediate Representation,多级中间表示)是一套编译器基础设施,不是可直接拿来加载模型并推理的终端运行时。
- 它用于在不同抽象层次表示与优化程序,许多编译器项目会在内部使用它。
IREE
- IREE 是基于 MLIR 的机器学习编译器与运行时,可将模型编译到多种硬件后端。
- 更适合有跨硬件编译、定制运行时或较深工程优化需求的团队;对于刚开始部署常规模型的学习者,ONNX Runtime、OpenVINO 或平台官方工具通常更容易上手。
Apache TVM
- TVM 是另一套开源机器学习编译器栈,能够针对不同硬件生成与优化算子实现。
- 它适合研究编译优化或需要适配特殊硬件的场景;使用时需要投入更多时间处理模型导入、算子支持、调优与部署集成。
9.2.3.5 总结与选型建议
没有“所有场景都最好的推理引擎”。可以先按“模型类型、目标设备、是否跨框架、性能目标和团队维护成本”缩小范围,再用真实模型测试确认。
| 典型需求 | 可优先了解的方案 | 选择原因 |
|---|---|---|
| 常规 CV、语音或传统机器学习模型,需要跨框架/跨语言部署 | ONNX + ONNX Runtime | 格式与运行时生态成熟,支持 CPU、GPU 与多种 Execution Provider。 |
| NVIDIA GPU 上追求固定模型的低延迟、高吞吐 | ONNX + TensorRT 或 TensorRT | 能针对 NVIDIA GPU 构建高度优化的推理 engine。 |
| 仅在 PyTorch/Python 服务内推理,先快速尝试优化 | torch.compile |
改动较小,便于直接在现有 PyTorch 代码中测试收益。 |
| Apple 的 iOS / macOS 端侧应用 | Core ML | 与 Apple 平台、系统 API 和硬件加速路径结合紧密。 |
| Android、嵌入式或资源受限端侧设备 | LiteRT 或 ExecuTorch | 都面向端侧;前者更偏 TensorFlow 生态,后者更贴近 PyTorch 工作流。 |
| Intel 硬件为主 | OpenVINO | 面向 Intel CPU、GPU、NPU 的部署与优化工具链。 |
| 本地、离线、量化 LLM | GGUF + llama.cpp |
模型文件与运行时生态对量化、本地资源占用有较多针对性优化。 |
| LLM 在线服务、高并发生成 | vLLM、SGLang、TensorRT-LLM | 除模型计算外,还会专门优化批处理调度与 KV Cache。 |
对本教程当前的学习阶段而言,可以先掌握一条最通用的主线:
1 | |
当明确知道目标设备或模型类型后,再选择 Core ML、OpenVINO、TensorRT、ExecuTorch 或 LLM 专用运行时,会比一开始就在众多工具之间切换更高效。