📚 课程目录第4章 PyTorch模型导出基础

4、PyTorch模型导出基础：使用torch.onnx.export导出简单模型

在完成环境准备与ONNX Runtime的初步验证后，我们正式进入模型转换的核心环节。本章将聚焦于PyTorch官方提供的torch.onnx.export函数，通过导出线性回归与小型CNN两个典型模型，深入理解其关键参数的含义与作用。这不仅是后续处理动态输入、自定义算子、量化部署等高级主题的基石，更是连接训练框架与推理引擎的“桥梁工程”。我们将从API原理出发，结合实战案例，剖析导出过程中的常见陷阱，帮助你建立一套可复用的导出方法论。

上图展示了从PyTorch模型到ONNX文件的核心转换链路。整个过程看似简单，但其中涉及模型追踪（Tracing）、算子映射、静态图构建等关键机制。理解这些底层原理，是正确设置参数、排查导出问题的前提。

4.1 核心原理：追踪（Tracing）机制与静态图构建

torch.onnx.export的核心机制是追踪（Tracing）。它并非直接读取PyTorch模型的源码逻辑，而是通过提供一个示例输入张量，让模型真实地执行一次前向传播。在这个过程中，PyTorch会记录下所有张量操作的执行顺序，并将这些动态执行的操作“固化”为一个静态的计算图，最后将这个计算图映射为ONNX的算子格式。

这种机制带来两个重要推论：

数据依赖控制流会被“拍平”：例如，模型中的if语句或for循环如果依赖于输入张量的值（如if x.sum() > 0），在追踪时只会记录当前输入所走的那一条分支。导出的模型将失去动态决策能力。
输入的形状（Shape）是固定的：默认情况下，导出的模型输入形状与示例输入完全一致。如果后续推理时输入尺寸不同，会直接报错。这需要通过dynamic_axes参数来声明动态维度。

  关键认知：ONNX模型是“数据流图”而非“程序”。它描述的是张量如何经过一系列算子流转到输出，不包含Python的控制流逻辑。因此，导出前必须确保模型内部没有依赖于数据值的控制流，或者使用torch.onnx.is_constant_trace等工具进行预检。

4.2 操作流程：torch.onnx.export 参数详解与导出步骤

我们以一个简单的线性回归模型为例，逐步拆解导出过程。首先定义模型并初始化：

import torch
import torch.nn as nn

class LinearRegression(nn.Module):
    def __init__(self, in_features=10, out_features=1):
        super().__init__()
        self.linear = nn.Linear(in_features, out_features)
    
    def forward(self, x):
        return self.linear(x)

# 实例化模型并设置为评估模式
model = LinearRegression()
model.eval()  # 重要：切换为eval模式，关闭Dropout/BatchNorm的随机行为

# 创建示例输入，batch_size=4, 特征维度=10
dummy_input = torch.randn(4, 10)

接下来，调用torch.onnx.export进行导出。以下代码展示了最核心的参数配置：

torch.onnx.export(
    model,                      # 1. 要导出的模型
    dummy_input,                # 2. 示例输入
    "linear_regression.onnx",   # 3. 输出文件路径
    export_params=True,         # 4. 是否导出权重参数
    opset_version=13,           # 5. ONNX算子集版本
    do_constant_folding=True,   # 6. 常量折叠优化
    input_names=['input'],      # 7. 输入张量名称
    output_names=['output'],    # 8. 输出张量名称
    dynamic_axes={              # 9. 动态维度声明
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)

下面对上述参数进行逐一解读：

参数名	作用	典型值 / 注意事项
`model`	要导出的PyTorch模型实例	必须已调用`.eval()`，且处于CPU上
`args`	示例输入张量（或元组）	用于触发前向传播，形状需符合模型预期
`f`	输出文件路径	建议以`.onnx`结尾
`export_params`	是否将模型权重参数固化到ONNX文件中	部署时通常设为`True`
`opset_version`	指定ONNX算子集版本	版本越高支持的新算子越多，但需考虑推理引擎兼容性。边缘端建议11~13
`do_constant_folding`	是否执行常量折叠优化	将固定计算（如权重归一化）预先算好，减小模型体积
`input_names / output_names`	为输入输出张量指定名称	便于后续在推理引擎中按名称绑定数据
`dynamic_axes`	声明哪些维度是动态的	格式为`{名称: {维度索引: 维度别名}}`，如`{'input': {0: 'batch'}}`

4.2.1 导出后的标准验证流程

导出成功后，必须进行两项基础验证：

结构完整性检查：使用onnx.checker.check_model验证ONNX文件是否符合规范。
数值一致性对比：使用ONNX Runtime加载模型，输入相同数据，对比PyTorch模型与ONNX模型的输出差异。

import onnx
import onnxruntime as ort
import numpy as np

# 1. 结构检查
onnx_model = onnx.load("linear_regression.onnx")
onnx.checker.check_model(onnx_model)
print("✓ ONNX模型结构检查通过")

# 2. 数值对比
# PyTorch推理
with torch.no_grad():
    torch_output = model(dummy_input).numpy()

# ONNX Runtime推理
ort_session = ort.InferenceSession("linear_regression.onnx")
ort_inputs = {ort_session.get_inputs()[0].name: dummy_input.numpy()}
ort_output = ort_session.run(None, ort_inputs)[0]

# 对比误差
np.testing.assert_allclose(torch_output, ort_output, rtol=1e-4, atol=1e-5)
print("✓ 数值一致性验证通过，最大误差:", np.abs(torch_output - ort_output).max())

4.3 案例拆解：小型CNN模型的导出与动态轴配置

线性回归模型过于简单，无法体现CNN中常见的多维张量操作。下面我们构建一个用于MNIST分类的小型CNN，并重点演示dynamic_axes在批处理维度上的应用。

class SmallCNN(nn.Module):
    def __init__(self, num_classes=10):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(1, 16, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(16, 32, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2)
        )
        self.classifier = nn.Sequential(
            nn.Flatten(),
            nn.Linear(32 * 7 * 7, 128),
            nn.ReLU(),
            nn.Linear(128, num_classes)
        )
    
    def forward(self, x):
        return self.classifier(self.features(x))

# 初始化模型
cnn_model = SmallCNN()
cnn_model.eval()

# 示例输入：batch=8, 单通道, 28x28
dummy_cnn_input = torch.randn(8, 1, 28, 28)

# 导出，声明batch维度为动态
torch.onnx.export(
    cnn_model,
    dummy_cnn_input,
    "small_cnn.onnx",
    input_names=['image_input'],
    output_names=['logits'],
    dynamic_axes={
        'image_input': {0: 'batch'},
        'logits': {0: 'batch'}
    },
    opset_version=13
)
print("✓ CNN模型导出完成")

导出后，我们可以在ONNX Runtime中测试不同的batch大小，验证动态轴是否生效：

ort_session = ort.InferenceSession("small_cnn.onnx")

# 测试 batch=1
test_input_1 = np.random.randn(1, 1, 28, 28).astype(np.float32)
output_1 = ort_session.run(None, {'image_input': test_input_1})[0]
print(f"Batch=1 输出形状: {output_1.shape}")

# 测试 batch=16
test_input_16 = np.random.randn(16, 1, 28, 28).astype(np.float32)
output_16 = ort_session.run(None, {'image_input': test_input_16})[0]
print(f"Batch=16 输出形状: {output_16.shape}")

经验提示：在边缘部署场景中，通常将batch维度设为动态，而将图像的高、宽设为固定值。因为边缘设备的输入分辨率往往是固定的（如320x240），而批处理大小可能因应用场景变化。此外，动态维度越多，ONNX Runtime的优化空间越小，推理性能可能下降，因此只对必要的维度声明动态。

4.4 常见误区与总结复盘

在导出过程中，初学者常会遇到以下几类问题，我们将其归纳为“三大误区”：

误区类型	具体表现	后果	正确做法
模型模式错误	忘记调用`model.eval()`	BatchNorm层使用批统计量，Dropout随机失活，导致导出图与训练行为不一致	导出前必须显式调用`model.eval()`
动态轴遗漏	未设置`dynamic_axes`	模型只能接受与示例输入完全相同的形状，部署时输入尺寸稍有变化即报错	根据部署需求，明确声明需要动态的维度
算子版本不匹配	`opset_version`设置过高	边缘端推理引擎（如TensorRT、OpenVINO）可能不支持最新算子，导致加载失败	查阅目标推理引擎支持的ONNX算子集版本，通常选择11~13

总结复盘：本章通过线性回归与小型CNN两个案例，完整演示了torch.onnx.export的标准用法。核心要点可归纳为：

理解追踪机制：导出是“执行一次”并记录操作，而非解析源码。
掌握关键参数：opset_version决定算子兼容性，dynamic_axes决定输入灵活性，export_params决定权重是否固化。
建立验证习惯：每次导出后，必须用onnx.checker和ONNX Runtime做双重验证，确保模型结构与数值均正确。

在下一阶段，我们将深入探讨如何处理更复杂的模型结构，例如包含LSTM、注意力机制或自定义操作的模型，这些场景将引入torch.onnx.export的进阶参数与自定义算子映射技术。