PyTorch模型基础：张量操作、自动求导、nn.Module、模型保存与加载

在将PyTorch模型导出为ONNX并最终转换到MLIR之前，我们必须先夯实PyTorch模型开发的核心基础。ONNX导出本质上是对模型计算图的静态捕获，这意味着模型中的每一个张量操作、参数管理方式以及控制流结构，都会直接影响导出结果的正确性与效率。本章将系统梳理张量操作、自动求导机制、nn.Module组织方式以及模型序列化方法，为后续的ONNX导出与MLIR转换建立坚实的知识底座。

核心原理：张量操作与计算图的关系

PyTorch的张量（Tensor）是构建神经网络的基本数据单元，其操作方式直接决定了计算图的拓扑结构。在ONNX导出场景下，我们需要理解张量操作的三个关键维度：形状语义、数据类型和设备位置。形状语义决定了ONNX图中的维度信息是否静态可推断；数据类型影响算子选择（如FP32与INT8对应不同的ONNX算子）；设备位置则关系到导出时是否需要处理设备间的数据拷贝节点。

计算图的构建是自动求导与ONNX导出的共同基础。当执行张量操作时，PyTorch会动态构建一个由Function节点连接的有向无环图（DAG）。每个节点保存了操作的输入输出张量以及梯度传播所需的中间变量。ONNX导出本质上是对这个动态图进行静态化快照——它遍历当前已构建的计算图，将PyTorch算子映射为ONNX算子，并记录张量的形状与数据类型信息。

张量操作实战：从创建到变形

在模型开发中，张量操作贯穿始终。以下代码展示了ONNX导出前需要重点关注的张量操作模式：

import torch
import torch.nn as nn

# 1. 创建与初始化
x = torch.randn(4, 3, 224, 224, dtype=torch.float32)  # 静态形状
batch_size = 8
y = torch.randn(batch_size, 3, 224, 224)  # 动态维度（batch_size可变）

# 2. 形状操作（影响ONNX静态形状推断）
z = x.view(4, -1)          # view要求内存连续
z2 = x.reshape(4, -1)      # reshape自动处理非连续
z3 = x.permute(0, 2, 3, 1) # 维度重排，产生非连续张量
z4 = x.transpose(1, 3)     # 转置操作

# 3. 广播与合并
a = torch.randn(4, 1, 224, 224)
b = torch.randn(1, 3, 224, 224)
c = a + b  # 广播机制，ONNX中对应MultidirectionalBroadcast

# 4. 索引与切片（ONNX导出时需谨慎）
d = x[:, 0:2, :, :]  # 切片操作在ONNX中可能产生Slice算子
e = x[torch.tensor([0, 2])]  # 高级索引，ONNX支持有限

# 5. 数据类型转换
f = x.to(torch.float16)  # 低精度转换，ONNX需匹配算子精度

在ONNX导出时，view、reshape、permute等形状操作会映射为Reshape、Transpose等ONNX算子。这些算子对输入形状的静态性要求较高——如果形状信息在导出时无法确定（例如依赖运行时计算的维度），ONNX图将包含动态维度，这会增加后续MLIR转换的复杂度。

自动求导机制：训练与推理的切换

自动求导是PyTorch训练的核心，但在ONNX导出时必须完全关闭。自动求导机制通过requires_grad标记和计算图记录实现梯度传播。当requires_grad=True时，每次前向传播都会构建梯度图，这会带来额外的内存开销和计算延迟。更重要的是，ONNX导出时如果梯度图仍然活跃，会导致导出器捕获到不必要的梯度节点，产生错误的ONNX图。

正确的做法是在导出前使用torch.no_grad()上下文管理器或调用model.eval()模式。这里需要区分两个概念：model.eval()影响的是BatchNorm和Dropout等层的行为（使用运行均值而非批统计量），而torch.no_grad()则完全禁用梯度追踪。两者在ONNX导出时都必不可少。

# 正确的导出前准备
model = MyModel()
model.load_state_dict(torch.load("model.pth"))
model.eval()  # 切换到推理模式

# 方式一：使用torch.no_grad()
with torch.no_grad():
    torch.onnx.export(model, dummy_input, "model.onnx")

# 方式二：全局禁用梯度
torch.set_grad_enabled(False)
torch.onnx.export(model, dummy_input, "model.onnx")

nn.Module：模型结构的组织与序列化

nn.Module是PyTorch模型的基本构建单元，它提供了参数管理、子模块注册和状态字典（state_dict）等核心功能。在ONNX导出场景下，nn.Module的组织方式直接影响导出图的清晰度和可转换性。

一个规范的nn.Module子类应包含：__init__方法中定义所有子模块和参数，forward方法中定义前向传播逻辑。PyTorch通过_modules字典自动追踪子模块，通过_parameters字典追踪可训练参数。这种结构化的管理方式使得state_dict能够完整保存模型的所有权重和偏置。

class ConvBlock(nn.Module):
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.conv = nn.Conv2d(in_ch, out_ch, 3, padding=1)
        self.bn = nn.BatchNorm2d(out_ch)
        self.relu = nn.ReLU(inplace=True)
        
    def forward(self, x):
        return self.relu(self.bn(self.conv(x)))

class MyModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.features = nn.Sequential(
            ConvBlock(3, 64),
            ConvBlock(64, 128),
            nn.MaxPool2d(2)
        )
        self.classifier = nn.Linear(128 * 112 * 112, 10)
        
    def forward(self, x):
        x = self.features(x)
        x = x.view(x.size(0), -1)
        return self.classifier(x)

在ONNX导出时，nn.Module的嵌套结构会被扁平化为ONNX图节点。子模块的边界不会保留在ONNX图中，所有操作都被展开为基本的ONNX算子。因此，模块划分的粒度不会影响ONNX图的结构，但会影响state_dict的组织和模型的可维护性。

模型保存与加载：state_dict、script与trace

PyTorch提供了多种模型序列化方式，每种方式在ONNX导出场景下有不同的适用性。理解这些差异是确保导出成功的关键。

state_dict方式是最轻量的保存方法，仅保存模型的参数和缓冲区（如BatchNorm的running_mean和running_var）。这种方式不包含模型结构信息，加载时需要先实例化模型再加载权重。对于ONNX导出，state_dict通常作为中间步骤——先加载权重到模型中，再执行导出。

torch.jit.script和torch.jit.trace是两种将PyTorch模型转换为TorchScript的方式。script通过解析Python源码生成完整的计算图，支持控制流（if/else、for循环）；trace通过实际执行模型记录操作序列，不支持数据依赖的控制流。在ONNX导出时，TorchScript模型可以直接作为输入，且script方式能更好地处理动态控制流。

案例拆解：从PyTorch模型到ONNX导出的完整流程

下面通过一个实际案例展示如何准备一个PyTorch模型并成功导出为ONNX。假设我们有一个用于图像分类的ResNet-18模型，需要导出到ONNX以便后续转换为MLIR。

背景与约束：模型在GPU上训练完成，输入为3×224×224的RGB图像，batch size固定为1。目标平台是边缘设备，要求FP32精度，且需要支持动态batch size（1~8）。

初始现象：直接使用torch.onnx.export导出时，出现"Failed to export"错误，提示某些算子不支持动态形状。

分析过程：检查模型代码发现，在forward函数中使用了x.view(x.size(0), -1)，其中x.size(0)是动态的。ONNX导出器无法推断view操作的输出形状，导致导出失败。

解决步骤：

1. 将x.view(x.size(0), -1)改为x.reshape(x.size(0), -1)，reshape在ONNX中映射为Reshape算子，支持动态维度。
2. 使用torch.jit.trace记录模型，输入使用固定batch size=1的dummy tensor。
3. 在导出时设置dynamic_axes参数，指定batch维度为动态：

import torch
import torchvision.models as models

model = models.resnet18(pretrained=True)
model.eval()

# 准备dummy输入
dummy_input = torch.randn(1, 3, 224, 224)

# 使用trace方式导出
traced_model = torch.jit.trace(model, dummy_input)

# 导出ONNX，支持动态batch
torch.onnx.export(
    traced_model,
    dummy_input,
    "resnet18.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={
        "input": {0: "batch_size"},
        "output": {0: "batch_size"}
    },
    opset_version=13
)

结果对比：导出成功后，使用onnxruntime验证模型推理结果与PyTorch原始模型一致（误差<1e-5）。动态batch支持验证：分别输入batch=1、4、8的tensor，均能正确推理。

风险与代价：动态batch会增加ONNX图的复杂度，可能导致部分算子无法完全静态优化。在MLIR转换阶段，动态维度需要转换为tensor.dim操作，增加了转换的复杂性。如果目标平台固定batch size，建议使用静态形状以获得最佳性能。

常见误区与总结复盘

在PyTorch模型准备阶段，开发者常犯以下错误，这些错误会在ONNX导出时集中爆发：

• 忽略eval模式：BatchNorm和Dropout在训练/推理模式下行为不同，未切换eval模式会导致导出的ONNX模型精度异常。
• 使用inplace操作：如nn.ReLU(inplace=True)，虽然节省内存，但可能干扰ONNX导出器的图捕获。
• 依赖Python控制流：在forward中使用if/else或for循环，且条件依赖输入数据，trace方式无法正确处理。
• 忽略设备一致性：模型参数在GPU上，而dummy_input在CPU上，导致导出时设备不匹配错误。
• 未检查算子兼容性：某些自定义操作或较新的PyTorch算子可能没有对应的ONNX映射。

总结本章核心要点：张量操作是计算图的基础，形状静态性决定ONNX导出的难易程度；自动求导必须在导出前完全关闭；nn.Module的结构化管理是state_dict和TorchScript的前提；选择合适的序列化方式（state_dict、script、trace）取决于模型是否包含动态控制流。掌握这些基础后，我们才能顺利进入ONNX导出的具体实践环节。