6.4 ONNX 与模型转换 (ONNX & Model Conversion)

文档摘要

6.4 ONNX 与模型转换 (ONNX & Model Conversion) 第六章：PyTorch 生态与工具 - 6.4 ONNX 与模型转换 (ONNX & Model Conversion) 在深度学习模型日益普及的今天，模型的可移植性和互操作性变得至关重要。不同的深度学习框架，如 PyTorch、TensorFlow、Caffe 等，各有优势，但同时也带来了模型部署和框架切换的挑战。为了解决这些问题，开放神经网络交换 (Open Neural Network Exchange, ONNX) 应运而生。ONNX 是一种开放的标准，旨在表示机器学习模型，使得模型可以在不同的框架之间进行转换和部署。

6.4 ONNX 与模型转换 (ONNX & Model Conversion)

第六章：PyTorch 生态与工具 - 6.4 ONNX 与模型转换 (ONNX & Model Conversion)

在深度学习模型日益普及的今天，模型的可移植性和互操作性变得至关重要。不同的深度学习框架，如 PyTorch、TensorFlow、Caffe 等，各有优势，但同时也带来了模型部署和框架切换的挑战。为了解决这些问题，开放神经网络交换 (Open Neural Network Exchange, ONNX) 应运而生。ONNX 是一种开放的标准，旨在表示机器学习模型，使得模型可以在不同的框架之间进行转换和部署。

本节将深入探讨 ONNX 以及在 PyTorch 中如何进行模型转换，包括 ONNX 的基本概念、PyTorch 中 ONNX 的导出与使用、模型检查与优化，以及 ONNX 在实际应用中的价值。

6.4.1 ONNX 简介：深度学习模型的通用语言

什么是 ONNX？

ONNX (Open Neural Network Exchange) 是一种用于表示机器学习模型的开放标准格式。它由 Microsoft、Facebook 和亚马逊等公司共同开发和维护，旨在实现不同深度学习框架之间的互操作性。ONNX 定义了一组通用的运算符集合和一种可扩展的文件格式，允许开发者在不同的框架之间轻松地迁移和部署模型。

为什么需要 ONNX？

框架互操作性 (Framework Interoperability): 深度学习领域存在多种优秀的框架，例如 PyTorch 以其灵活性和易用性著称，TensorFlow 则在工业界应用广泛。ONNX 允许模型在这些框架之间自由转换，避免了被特定框架绑定的风险，并可以充分利用各个框架的优势。
硬件加速和部署优化 (Hardware Acceleration and Deployment Optimization): 许多硬件厂商和推理引擎 (例如 NVIDIA TensorRT, Intel OpenVINO, Qualcomm SNPE) 都支持 ONNX 格式。通过将模型转换为 ONNX，可以更容易地利用这些硬件加速库和优化工具，提升模型在不同平台上的推理性能。
模型共享与合作 (Model Sharing and Collaboration): ONNX 提供了一种标准化的模型交换格式，方便研究人员和开发者共享和复用模型，促进了深度学习社区的合作与发展。

ONNX 的工作流程

ONNX 的核心工作流程可以概括为以下几个步骤：

模型导出为 ONNX： 使用训练框架 (例如 PyTorch) 的 ONNX 导出工具，将训练好的模型转换为 ONNX 格式。
ONNX 模型优化 (可选)： 使用 ONNX 提供的优化工具，对导出的 ONNX 模型进行图优化、节点融合等操作，以提升推理性能。
推理引擎加载 ONNX 模型： 使用支持 ONNX 格式的推理引擎 (例如 ONNX Runtime, TensorRT) 加载 ONNX 模型。
在目标硬件平台执行推理： 推理引擎将 ONNX 模型部署到目标硬件平台 (例如 CPU, GPU, 移动设备) 上进行推理。

6.4.2 PyTorch 中的 ONNX 导出：`torch.onnx.export`

PyTorch 提供了强大的 ONNX 导出功能，通过 torch.onnx.export 函数，可以将 PyTorch 模型转换为 ONNX 格式。

torch.onnx.export 函数详解

torch.onnx.export 函数的基本用法如下：


torch.onnx.export(model,                                # 要导出的 PyTorch 模型
                  args,                                 # 模型的输入，可以是 tuple 或 Tensor
                  f,                                    # 输出 ONNX 模型的路径 (文件路径或文件对象)
                  export_params=True,                     # 是否导出模型的参数 (权重)
                  opset_version=None,                     # ONNX 算子集版本，默认为 ONNX 稳定版本
                  do_constant_folding=True,               # 是否执行常量折叠优化
                  input_names=None,                       # 输入节点的名称列表
                  output_names=None,                      # 输出节点的名称列表
                  dynamic_axes=None,                      # 定义动态轴，用于处理变长输入/输出
                  verbose=False,                          # 是否打印详细导出信息
                  training=torch.onnx.TrainingMode.EVAL,  # 导出模型的状态 (训练或推理)
                  **kwargs)

常用参数解释：

model (nn.Module): 需要导出的 PyTorch 模型实例。
args (tuple 或 Tensor): 模型的输入，用于追踪模型的前向计算图。通常需要提供一个或多个与模型输入形状和类型匹配的 Tensor 或 tuple。
f (str 或 file object): 导出的 ONNX 模型的保存路径，可以是文件路径字符串，也可以是文件对象。
export_params (bool, default=True): 控制是否将模型的参数 (权重和偏置) 导出到 ONNX 文件中。通常设置为 True，以便 ONNX 模型包含完整的模型定义和权重信息。
opset_version (int, optional): 指定导出的 ONNX 算子集版本。ONNX 算子集版本不断演进，新版本通常会引入新的算子和功能。为了保证模型的兼容性，建议根据目标推理引擎的要求选择合适的 opset_version。常用的版本包括 9, 10, 11, 12, 13, 14, 15 等。
do_constant_folding (bool, default=True): 是否执行常量折叠优化。常量折叠是一种图优化技术，可以将计算图中可以预先计算的常量节点进行计算，并用结果替换这些节点，从而简化计算图，提升推理性能。
input_names (list of str, optional): 为模型的输入节点指定名称。这对于后续的模型检查和推理过程中的输入指定很有帮助。
output_names (list of str, optional): 为模型的输出节点指定名称。与 input_names 类似，方便模型检查和推理输出处理。
dynamic_axes (dict, optional): 用于定义模型输入和输出中的动态轴。动态轴表示该维度的大小可以在推理时发生变化，例如批次大小、序列长度等。这对于处理变长输入 (例如文本、音频) 的模型非常重要。dynamic_axes 是一个字典，键可以是 input_names 或 output_names 中的名称，值可以是维度索引列表或维度名称列表。
training (torch.onnx.TrainingMode, default=torch.onnx.TrainingMode.EVAL): 指定导出的模型状态。torch.onnx.TrainingMode.EVAL 表示导出推理模式的模型，torch.onnx.TrainingMode.TRAINING 表示导出训练模式的模型 (通常用于 ONNX 的训练功能，较少使用)。

代码实践：导出简单的线性回归模型


import torch
import torch.nn as nn
# 1. 定义一个简单的线性回归模型
class LinearRegression(nn.Module):
    def __init__(self):
        super(LinearRegression, self).__init__()
        self.linear = nn.Linear(1, 1) # 单输入单输出线性层
    def forward(self, x):
        return self.linear(x)
model = LinearRegression()
# 2. 创建一个虚拟输入
dummy_input = torch.randn(1, 1) # 形状为 (1, 1) 的随机张量
# 3. 指定 ONNX 模型的输出路径
onnx_path = "linear_regression.onnx"
# 4. 使用 torch.onnx.export 导出模型
torch.onnx.export(model,                     # 要导出的模型
                  dummy_input,               # 模型的虚拟输入
                  onnx_path,                 # 输出路径
                  verbose=True,                # 打印导出信息
                  input_names=['input'],      # 输入节点名称
                  output_names=['output'])    # 输出节点名称
print(f"ONNX 模型已导出到: {onnx_path}")

代码详解：

定义模型： 我们首先定义了一个简单的线性回归模型 LinearRegression，它包含一个线性层 nn.Linear(1, 1)。
创建虚拟输入： 为了让 torch.onnx.export 函数能够追踪模型的前向计算图，我们需要提供一个虚拟输入 dummy_input。这里我们创建了一个形状为 (1, 1) 的随机张量，模拟模型的输入。
指定输出路径： 我们指定导出的 ONNX 模型保存路径为 linear_regression.onnx。
导出模型： 调用 torch.onnx.export 函数，传入模型 model、虚拟输入 dummy_input、输出路径 onnx_path，并设置 verbose=True 以打印导出信息，并指定了输入和输出节点的名称分别为 'input' 和 'output'。

运行这段代码后，将在当前目录下生成 linear_regression.onnx 文件，这就是导出的 ONNX 模型。控制台会输出详细的导出信息，包括 ONNX 模型的结构、算子等。

6.4.3 ONNX 模型检查与可视化：Netron 和 ONNX Runtime 工具

导出 ONNX 模型后，我们需要对模型进行检查，确保导出的模型结构和算子符合预期。常用的 ONNX 模型检查和可视化工具包括 Netron 和 ONNX Runtime 提供的工具。

1. Netron：可视化 ONNX 模型

Netron 是一款开源的模型可视化工具，支持多种模型格式，包括 ONNX。Netron 可以直观地展示 ONNX 模型的计算图结构、节点信息、属性和参数。

使用 Netron 查看 linear_regression.onnx 模型：

安装 Netron： 可以从 Netron 官网下载对应平台的安装包，或者使用 pip 安装：pip install netron
启动 Netron： 在命令行中输入 netron 启动 Netron 应用。
加载 ONNX 模型： 在 Netron 应用中，选择 "File" -> "Open"，打开 linear_regression.onnx 文件。

Netron 将会加载并可视化 ONNX 模型。你可以看到模型的输入节点 'input'，线性层节点 'Gemm' (Generalized Matrix Multiplication，线性层的底层实现)，以及输出节点 'output'。点击节点可以查看节点的属性和参数信息。

2. ONNX Runtime 工具：代码检查 ONNX 模型

ONNX Runtime (ORT) 是一个高性能的推理引擎，同时也提供了一些用于 ONNX 模型检查的工具，例如 onnx.checker.check_model 和 onnx.shape_inference.infer_shapes。

代码实践：使用 ONNX Runtime 工具检查 ONNX 模型


import onnx
from onnx.checker import check_model
from onnx.shape_inference import infer_shapes
onnx_model_path = "linear_regression.onnx"
# 1. 加载 ONNX 模型
onnx_model = onnx.load(onnx_model_path)
# 2. 使用 onnx.checker 检查模型
try:
    check_model(onnx_model)
    print("ONNX 模型检查通过！")
except onnx.checker.ValidationError as e:
    print(f"ONNX 模型检查失败：\n{e}")
# 3. 使用 onnx.shape_inference 推理形状 (可选)
try:
    inferred_model = infer_shapes(onnx_model)
    print("ONNX 模型形状推理成功！")
    # 可以进一步访问 inferred_model.graph.value_info 获取形状信息
except Exception as e:
    print(f"ONNX 模型形状推理失败：\n{e}")

代码详解：

加载 ONNX 模型： 使用 onnx.load 函数加载 ONNX 模型文件。
模型检查： 使用 onnx.checker.check_model 函数对 ONNX 模型进行结构和语义检查，例如检查算子是否合法、输入输出类型是否匹配等。如果模型通过检查，会打印 "ONNX 模型检查通过！"，否则会打印错误信息。
形状推理 (可选)： 使用 onnx.shape_inference.infer_shapes 函数尝试推理 ONNX 模型中各个节点的输出形状。形状推理可以帮助我们验证模型的形状信息是否正确，并为后续的优化和推理提供信息。

6.4.4 ONNX 模型推理：ONNX Runtime

ONNX Runtime (ORT) 是一个跨平台的推理引擎，可以高效地执行 ONNX 模型。ORT 支持多种硬件平台 (CPU, GPU, 移动设备) 和操作系统，并提供了丰富的 API (Python, C++, C#, Java 等)。

代码实践：使用 ONNX Runtime 推理 ONNX 模型


import onnxruntime as ort
import numpy as np
onnx_model_path = "linear_regression.onnx"
# 1. 创建 ONNX Runtime 推理会话
ort_session = ort.InferenceSession(onnx_model_path)
# 2. 获取输入和输出节点信息
input_name = ort_session.get_inputs()[0].name
output_name = ort_session.get_outputs()[0].name
print(f"输入节点名称: {input_name}")
print(f"输出节点名称: {output_name}")
# 3. 准备输入数据 (需要转换为 numpy 数组)
input_data = np.array([[2.0]], dtype=np.float32) # 输入值 2.0，形状 (1, 1)
# 4. 执行推理
ort_inputs = {input_name: input_data}
ort_outputs = ort_session.run([output_name], ort_inputs)
# 5. 获取推理结果
output_result = ort_outputs[0]
print(f"推理结果: {output_result}")

代码详解：

创建推理会话： 使用 ort.InferenceSession(onnx_model_path) 创建 ONNX Runtime 推理会话，加载 ONNX 模型。
获取输入/输出节点信息： 通过 ort_session.get_inputs() 和 ort_session.get_outputs() 获取模型的输入和输出节点信息，包括节点名称、形状、数据类型等。
准备输入数据： ONNX Runtime 的输入数据需要是 numpy 数组。我们将输入值 2.0 转换为形状为 (1, 1) 的 numpy float32 数组。
执行推理： 调用 ort_session.run([output_name], ort_inputs) 执行推理。第一个参数是需要获取的输出节点名称列表，第二个参数是输入数据字典，键为输入节点名称，值为输入数据 numpy 数组。
获取推理结果： ort_session.run 返回一个输出结果列表，我们取出第一个元素 ort_outputs[0]，即为模型的输出结果。

运行这段代码，ONNX Runtime 将加载 linear_regression.onnx 模型，并使用输入值 2.0 进行推理，输出线性回归模型的预测结果。

6.4.5 动态轴 (Dynamic Axes)：处理变长输入

在处理序列数据 (例如文本、音频) 或变长输入时，模型的输入形状可能不是固定的。ONNX 提供了动态轴 (Dynamic Axes) 的概念，允许模型在某些维度上接受变长的输入。在 PyTorch 中导出 ONNX 模型时，可以使用 dynamic_axes 参数来指定动态轴。

代码实践：导出带有动态轴的 LSTM 模型


import torch
import torch.nn as nn
# 1. 定义一个简单的 LSTM 模型
class LSTMModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size):
        super(LSTMModel, self).__init__()
        self.lstm = nn.LSTM(input_size, hidden_size, batch_first=True) # batch_first=True 表示输入形状为 (batch, seq_len, input_size)
        self.linear = nn.Linear(hidden_size, output_size)
    def forward(self, x):
        out, _ = self.lstm(x)
        out = self.linear(out[:, -1, :]) # 取最后一个时间步的输出
        return out
input_size = 10
hidden_size = 20
output_size = 5
model = LSTMModel(input_size, hidden_size, output_size)
# 2. 创建虚拟输入，批次大小和序列长度可变
dummy_input = torch.randn(2, 5, input_size) # 示例输入：batch_size=2, seq_len=5
# 3. 指定 ONNX 模型输出路径
onnx_path = "lstm_model_dynamic_axes.onnx"
# 4. 定义动态轴：批次大小 (batch_size) 和序列长度 (sequence)
dynamic_axes = {'input': {0: 'batch_size', 1: 'sequence'},    # 输入的第0维 (batch_size) 和第1维 (sequence) 是动态的
                'output': {0: 'batch_size'}}                   # 输出的第0维 (batch_size) 是动态的
# 5. 使用 torch.onnx.export 导出模型，并指定 dynamic_axes
torch.onnx.export(model,
                  dummy_input,
                  onnx_path,
                  verbose=True,
                  input_names=['input'],
                  output_names=['output'],
                  dynamic_axes=dynamic_axes)
print(f"带有动态轴的 LSTM 模型已导出到: {onnx_path}")

代码详解：

定义 LSTM 模型： 我们定义了一个简单的 LSTM 模型 LSTMModel，使用 nn.LSTM 层和 nn.Linear 层。batch_first=True 参数表示 LSTM 层的输入形状为 (batch, seq_len, input_size)。
创建虚拟输入： 我们创建了一个示例虚拟输入 dummy_input，形状为 (2, 5, input_size)，其中批次大小为 2，序列长度为 5。
定义动态轴： 我们使用 dynamic_axes 字典来定义动态轴。对于输入 'input'，我们将第 0 维 (批次大小) 和第 1 维 (序列长度) 指定为动态轴，并分别命名为 'batch_size' 和 'sequence'。对于输出 'output'，我们将第 0 维 (批次大小) 指定为动态轴。
导出模型： 在 torch.onnx.export 函数中，我们将 dynamic_axes 参数设置为我们定义的动态轴字典。

导出后，可以使用 Netron 查看 lstm_model_dynamic_axes.onnx 模型，可以看到输入和输出节点的形状信息中，动态轴维度被标记为可变的 (例如 batch_size, sequence)。在使用 ONNX Runtime 推理时，可以输入不同批次大小和序列长度的数据。

6.4.6 模型转换与优化：超越框架限制

ONNX 不仅仅是一个模型交换格式，它还是模型转换和优化的桥梁。通过将模型转换为 ONNX，我们可以利用各种 ONNX 工具和推理引擎，将模型部署到不同的硬件平台，并进行性能优化。

模型转换工具示例：

ONNX-TensorRT: 用于将 ONNX 模型转换为 NVIDIA TensorRT 引擎可执行的模型，以利用 NVIDIA GPU 的高性能推理能力。
ONNX-OpenVINO: 用于将 ONNX 模型转换为 Intel OpenVINO 工具包可执行的模型，以在 Intel CPU 和 GPU 上进行优化推理。
CoreMLTools: 用于将 ONNX 模型转换为 Apple CoreML 格式，以在 iOS 和 macOS 设备上进行硬件加速推理。
ONNX-TF: 用于将 ONNX 模型转换为 TensorFlow SavedModel 格式，实现 ONNX 到 TensorFlow 的转换。

模型优化工具示例：

ONNX Optimizer: ONNX 官方提供的模型优化工具，可以进行图优化 (例如常量折叠、节点融合、算子简化等)，减少模型大小，提升推理速度。
Quantization Tools: ONNX 生态系统中也包含一些量化工具，可以将 ONNX 模型量化为 INT8 或更低精度，进一步压缩模型大小，加速推理。

ONNX 模型转换和优化流程：

PyTorch 模型导出为 ONNX： 使用 torch.onnx.export 将 PyTorch 模型转换为 ONNX 格式。
ONNX 模型优化 (可选)： 使用 ONNX Optimizer 对 ONNX 模型进行图优化，提升推理性能。
使用模型转换工具转换为目标平台格式： 根据目标平台选择合适的模型转换工具 (例如 ONNX-TensorRT, ONNX-OpenVINO, CoreMLTools)，将 ONNX 模型转换为目标平台特定的模型格式。
在目标平台部署和推理： 使用目标平台提供的推理引擎加载转换后的模型，并在目标硬件平台上进行高性能推理。

6.4.7 总结与展望

ONNX 作为深度学习模型的通用交换格式，在模型的可移植性、互操作性和部署优化方面发挥着重要作用。PyTorch 提供了强大的 ONNX 导出功能，使得 PyTorch 模型可以轻松地转换为 ONNX 格式，并借助 ONNX 生态系统中的工具，实现跨框架、跨平台的模型部署和优化。

随着深度学习技术的不断发展，ONNX 标准也在持续演进，支持更多的算子、优化技术和硬件平台。掌握 ONNX 与模型转换技术，对于深度学习工程师来说至关重要，可以帮助我们更好地利用不同的框架和硬件资源，加速模型部署和应用落地。

未来，ONNX 将继续在深度学习领域扮演关键角色，推动深度学习模型的标准化和生态系统的繁荣发展。我们期待 ONNX 能够支持更多的新兴技术 (例如 Transformer, 量化, 剪枝)，并为深度学习的应用带来更广阔的空间。

作者与出处

原作者：灏天文库

来源：灏天文库

原文： https://aiknowledge.cn/collection/215-pytorch%E6%B7%B1%E5%BA%A6%E5%AD%A6%E4%B9%A0%E6%A1%86%E6%9E%B6%E5%AE%9E%E6%88%98

整理：灏天文库整理

由灏天文库平台收录，内容或由平台用户上传，仅供学习交流

6.4 ONNX 与模型转换 (ONNX & Model Conversion)

文档摘要