4、模型导出入门：使用torch.onnx.export和tf2onnx完成第一个简单模型的ONNX导出

在完成ONNX Runtime部署、算子兼容性检查与推理优化之前，我们首先需要解决一个最基础也最关键的问题：如何将训练好的PyTorch或TensorFlow模型转换为标准的ONNX格式。这一过程看似简单，实则涉及模型计算图的静态化、动态轴的声明、算子映射的兼容性以及输入输出的元信息定义等多个环节。本章将带领你从零开始，分别使用torch.onnx.export和tf2onnx工具，完成第一个简单模型的ONNX导出，并深入剖析导出过程中的核心原理与常见陷阱。

PyTorch 导出链路 ① 定义模型并设为 eval 模式 model.eval() + torch.no_grad() ② 构造 dummy 输入张量 torch.randn(1, 3, 224, 224) ③ 调用 torch.onnx.export 设置 input_names / output_names / dynamic_axes ④ 生成 model.onnx 文件 包含计算图、权重、算子版本信息 TensorFlow 导出链路 ① 保存/加载 Keras 模型 model.save('model.h5') 或 tf.saved_model ② 调用 tf2onnx.convert 指定输入/输出节点名或使用签名 ③ 设置 opset 与 target 版本 --target onnxruntime / --opset 13 ④ 生成 model.onnx 文件 使用 graph freeze 与优化 统一 ONNX 中间表示

上图清晰地展示了两种主流深度学习框架导出ONNX模型的整体链路。PyTorch通过torch.onnx.export直接对模型实例进行追踪（Tracing）或脚本化（Scripting）导出；TensorFlow则通过tf2onnx工具将SavedModel或Keras模型转换为ONNX格式。两条路径最终汇合到统一的ONNX中间表示，为后续的跨平台部署奠定基础。接下来，我们将分别深入探讨这两种导出方式的具体操作与核心原理。

4.1 核心原理：从动态图到静态图的转换

理解ONNX导出的本质，是掌握整个自动化流程的基石。无论是PyTorch还是TensorFlow，其训练时的计算图都是动态的、面向自动微分的；而ONNX作为一种推理格式，需要的是静态的、确定性的计算图。这种转换过程涉及几个关键的技术要点。

4.1.1 PyTorch的追踪（Tracing）与脚本化（Scripting）

PyTorch默认使用追踪（Tracing）机制进行导出。其原理是：构造一个虚拟输入（dummy input），让模型真实地执行一次前向传播，同时记录下所有执行过的张量操作，将这些操作序列化为一个静态的计算图。这种方式的优点是简单、自动化程度高，但缺点是无法捕获依赖数据控制流的逻辑，例如if语句、for循环等。如果模型内部包含这类动态控制流，追踪导出会将这些分支固定为追踪时执行的那一条路径，导致导出的模型行为与原始模型不一致。

为了解决这个问题，PyTorch提供了脚本化（Scripting）模式。通过torch.jit.script将模型转换为TorchScript，再导出为ONNX。Scripting会解析模型的Python源码，将其编译为静态图，从而保留控制流逻辑。但Scripting对模型代码的写法有较高要求，需要避免使用Python的动态特性。

核心要点：对于大多数标准的卷积神经网络（CNN）或Transformer模型，使用Tracing模式即可满足需求。如果模型包含if、for等控制流，且这些控制流依赖于输入张量的值，则必须使用Scripting或确保追踪时所有分支都被执行到。

4.1.2 TensorFlow的图冻结与节点映射

TensorFlow模型通常以SavedModel或Keras的H5格式保存。这些格式中包含了完整的训练图、变量、优化器状态等。tf2onnx的核心工作是将TensorFlow的计算图（GraphDef）转换为ONNX的计算图。这个过程包括：

  1. 图冻结（Freeze Graph）：将变量（Variable）转换为常量（Constant），去除训练相关的节点（如梯度计算、优化器更新），得到一个纯粹的推理图。
  2. 算子映射（Op Mapping）：将TensorFlow的算子（如Conv2D、MatMul、Relu）映射到ONNX标准算子。由于两个框架的算子集并非一一对应，tf2onnx内部维护了一个庞大的映射表，并会插入一些辅助节点（如Transpose、Reshape）来对齐数据布局和维度。
  3. 图优化（Graph Optimization）：在转换过程中，tf2onnx会执行一些常量折叠、冗余节点消除等优化操作，使生成的ONNX模型更加精简。

4.2 操作流程：PyTorch模型导出实战

下面我们通过一个具体的例子，演示如何使用torch.onnx.export导出一个简单的卷积神经网络模型。我们以经典的LeNet-5模型为例，该模型结构简单，适合作为入门演示。

4.2.1 环境准备与模型定义

首先，确保已安装PyTorch和ONNX。然后定义LeNet-5模型：

import torch
import torch.nn as nn
import torch.onnx

class LeNet5(nn.Module):
    def __init__(self):
        super(LeNet5, self).__init__()
        self.conv1 = nn.Conv2d(1, 6, kernel_size=5, padding=2)
        self.pool1 = nn.AvgPool2d(kernel_size=2, stride=2)
        self.conv2 = nn.Conv2d(6, 16, kernel_size=5)
        self.pool2 = nn.AvgPool2d(kernel_size=2, stride=2)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool1(torch.relu(self.conv1(x)))
        x = self.pool2(torch.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = self.fc3(x)
        return x

model = LeNet5()
model.eval()  # 切换到评估模式，关闭Dropout和BatchNorm的统计更新

4.2.2 构造Dummy输入并执行导出

接下来，构造一个形状为(1, 1, 28, 28)的随机张量作为示例输入，并调用导出函数：

dummy_input = torch.randn(1, 1, 28, 28)

torch.onnx.export(
    model,                      # 要导出的模型
    dummy_input,                # 示例输入张量
    "lenet5.onnx",              # 输出文件路径
    export_params=True,         # 是否导出权重参数
    opset_version=13,           # ONNX算子集版本
    do_constant_folding=True,   # 是否执行常量折叠优化
    input_names=['input'],      # 输入节点名称
    output_names=['output'],    # 输出节点名称
    dynamic_axes={
        'input': {0: 'batch_size'},   # 将第0维（batch）设为动态
        'output': {0: 'batch_size'}
    }
)
print("模型导出成功！")

在上述代码中，dynamic_axes参数非常关键。它允许我们指定哪些维度是动态的。例如，{0: 'batch_size'}表示输入和输出的第0维（即批量大小）是动态的，这样导出的模型就可以接受任意batch size的输入。如果不设置该参数，导出的模型将固定输入形状为(1, 1, 28, 28)。

经验提醒：在设置opset_version时，建议使用较新的版本（如13或更高），因为新版本支持更多的算子，且对动态形状的支持更好。但也要注意，过高的opset版本可能导致旧版ONNX Runtime无法加载。通常，ONNX Runtime 1.10及以上版本支持opset 13。

4.3 操作流程：TensorFlow模型导出实战

对于TensorFlow模型，我们使用tf2onnx工具进行转换。首先需要安装tf2onnx：pip install tf2onnx。然后，我们以一个简单的Keras Sequential模型为例。

4.3.1 创建并保存Keras模型

import tensorflow as tf
import tf2onnx

# 构建一个简单的Keras模型
model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(32, (3, 3), activation='relu', input_shape=(28, 28, 1)),
    tf.keras.layers.MaxPooling2D((2, 2)),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(10, activation='softmax')
])

# 保存为SavedModel格式（推荐）
model.save('tf_model', save_format='tf')
print("TensorFlow模型已保存")

4.3.2 使用tf2onnx命令行工具转换

tf2onnx提供了命令行工具和Python API两种方式。这里我们演示命令行方式，它更为直观：

python -m tf2onnx.convert \
    --saved-model tf_model \
    --output tf_model.onnx \
    --opset 13 \
    --target onnxruntime

参数说明：

  • --saved-model：指定输入的SavedModel目录。
  • --output：指定输出的ONNX文件路径。
  • --opset：指定ONNX算子集版本。
  • --target：指定目标运行时，这里选择onnxruntime，tf2onnx会针对该运行时进行特定优化。

如果使用Python API，可以这样写：

import tf2onnx

# 加载SavedModel
model = tf.saved_model.load('tf_model')
# 获取推理函数（signature）
spec = model.signatures['serving_default']
# 转换为ONNX
onnx_model, _ = tf2onnx.convert.from_function(
    spec,
    input_signature=[tf.TensorSpec([None, 28, 28, 1], tf.float32, name='input')],
    output_path='tf_model.onnx',
    opset=13
)
print("转换完成")

风险提示：在转换TensorFlow模型时，务必确保输入输出的张量形状和数据类型与训练时一致。特别是当模型包含预处理或后处理逻辑时，建议将这些逻辑移出模型，只转换核心推理部分，以避免不必要的算子映射问题。

4.4 案例拆解：导出后的验证与常见问题排查

导出ONNX模型只是第一步，验证导出的模型是否正确、能否被ONNX Runtime正确加载并产生与原始框架一致的推理结果，是更为关键的环节。下面我们通过一个完整的案例来演示验证流程。

4.4.1 使用ONNX Runtime进行推理对比

我们以刚才导出的PyTorch LeNet-5模型为例，使用ONNX Runtime加载并推理，与PyTorch原始模型的输出进行对比。

import onnxruntime as ort
import numpy as np

# 创建ONNX Runtime推理会话
ort_session = ort.InferenceSession('lenet5.onnx')

# 准备输入数据（与导出时的dummy_input相同分布）
test_input = torch.randn(1, 1, 28, 28)

# PyTorch推理
with torch.no_grad():
    torch_output = model(test_input).numpy()

# ONNX Runtime推理
ort_inputs = {ort_session.get_inputs()[0].name: test_input.numpy()}
ort_output = ort_session.run(None, ort_inputs)[0]

# 对比输出差异
diff = np.abs(torch_output - ort_output).max()
print(f"PyTorch与ONNX Runtime最大输出差异: {diff:.6e}")

# 设置一个阈值，判断是否一致
if diff < 1e-5:
    print("✅ 模型导出验证通过！")
else:
    print("❌ 模型导出存在较大差异，请检查模型结构或导出参数。")

通常，由于浮点数运算的微小差异，两个框架的输出不会完全一致，但最大差异应在1e-5量级以内。如果差异过大，则说明导出过程中出现了问题。

4.4.2 常见问题与排查策略

问题现象 可能原因 排查与解决方案
导出时报错 Unsupported operator 模型使用了ONNX不支持的算子，或opset版本过低 升级opset_version；检查是否有自定义算子，需要注册或替换；使用torch.onnx.is_onnx_backend_available()检查支持情况
导出的模型推理结果与原始框架不一致 模型包含动态控制流，Tracing模式未捕获所有分支；BatchNorm或Dropout未正确设置为eval模式 使用torch.jit.script进行Scripting导出；确保调用model.eval()；检查输入数据的预处理是否一致
ONNX Runtime加载失败，提示Invalid Graph 模型图中存在悬空节点或形状不匹配 使用onnx.checker.check_model检查模型完整性；使用onnx.shape_inference.infer_shapes进行形状推断
动态轴设置后，推理速度变慢 动态轴导致ONNX Runtime无法进行一些静态形状优化 如果业务场景中batch size固定，建议不要设置动态轴；或者设置多个固定形状的优化配置

4.4.3 案例复盘：一个BatchNorm导致的导出陷阱

我们曾遇到一个实际案例：某团队导出一个包含BatchNorm层的ResNet模型，在PyTorch中推理正常，但导出为ONNX后，在ONNX Runtime上的推理结果完全错误。经过排查，发现原因是在导出前没有调用model.eval()。在训练模式下，BatchNorm层会使用当前batch的均值和方差进行归一化，而导出时使用的是dummy input，其统计量是随机的，导致导出的模型权重中包含了错误的BatchNorm参数。

解决方案：在导出前显式调用model.eval()，将BatchNorm层切换到推理模式，使用训练阶段累积的全局统计量。这个案例提醒我们，导出前的模型状态设置是决定导出正确性的第一道关卡。

4.5 常见误区与总结复盘

在模型导出的入门阶段，有几个高频误区需要特别留意，它们往往会导致导出失败或推理结果异常。

误区一：忽略model.eval()与torch.no_grad()。这会导致BatchNorm和Dropout层行为异常，且计算图会包含梯度计算节点，使模型臃肿且行为错误。正确做法是：导出前必须调用model.eval()，并在推理验证时使用torch.no_grad()。

误区二：盲目追求高opset版本。虽然高opset支持更多算子，但可能导致模型无法在旧版ONNX Runtime上运行。建议根据部署环境选择合适的opset版本，并在导出后使用目标环境的ONNX Runtime进行验证。

误区三：忽视输入输出的命名与动态轴。如果导出的模型需要被其他服务调用，输入输出名称必须具有可读性且稳定。动态轴的设置要谨慎，过多动态维度会增加推理引擎的优化难度，影响性能。

总结复盘：完成第一个模型的ONNX导出，核心在于理解“动态图转静态图”的本质。对于PyTorch，要掌握Tracing与Scripting的适用场景；对于TensorFlow，要熟悉SavedModel的导出与tf2onnx的参数配置。导出后的验证环节不可或缺，通过对比原始框架与ONNX Runtime的输出差异，可以快速定位问题。建议在项目初期就建立一套自动化的导出与验证脚本，将模型导出纳入CI/CD流程，避免人工操作带来的不确定性。

← 上一章 📖 返回目录 下一章 →