📚 课程目录第2章 PyTorch 模型基础

2. PyTorch 模型基础

在把 PyTorch 模型转换到 ONNX、再进入 MLIR 生态之前，必须先理解 PyTorch 自身的模型组织方式与执行机制。很多导出失败、算子缺失、动态维度错乱的问题，根源并不在转换工具，而在于模型本身没有为“静态化导出”做好准备。本节围绕 nn.Module、动态图、TorchScript 三条主线，梳理导出前的准备工作与常见坑点。

一、核心原理：nn.Module 是模型的“结构契约”

nn.Module 是 PyTorch 中一切可训练组件的基类。它通过 __init__ 注册子模块与参数，通过 forward 定义前向计算。导出工具（torch.onnx.export、torch.jit.trace）本质上是在“追踪”或“编译”forward 中实际执行的操作序列，因此 forward 的写法直接决定了导出图的质量。

关键点在于：nn.Module 的 forward 应当只包含张量计算，不应混入 Python 控制流、打印、文件读写、随机采样等副作用操作。导出器无法把这些 Python 语义翻译成 ONNX 或 MLIR 算子，它们要么被静默忽略，要么直接报错。

导出前的第一原则：让 forward 成为纯函数式的张量计算图。所有分支、循环、动态形状，都要么在导出前固化为常量，要么用可被追踪的等价写法表达。

二、动态图 vs 静态图：追踪机制的本质

PyTorch 默认是动态图（Eager Mode）：每次前向都会重新构建计算图，Python 的 if、for、while 会真实执行。而 ONNX 和 MLIR 面向的是静态图，需要一次性确定算子拓扑。二者之间的桥梁就是“追踪（Tracing）”或“脚本化（Scripting）”。

追踪只记录某一次前向中实际走过的路径。如果模型里有 if x.sum() > 0 这类数据相关分支，追踪只会保留当次命中的那一支，另一支被丢弃，导致导出模型在别的输入下行为错误。这是最典型、也最隐蔽的导出事故。

三、TorchScript：把模型变成可序列化的 IR

TorchScript 是 PyTorch 自带的中间表示，它把 Python 模型编译成可脱离 Python 解释器运行的图。它有两种生成方式：trace 与 script。前者适合无数据相关控制流的模型，后者适合含 if、循环、List/Dict 等结构的模型。

import torch
import torch.nn as nn

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(8, 4)
        self.relu = nn.ReLU()

    def forward(self, x):
        return self.relu(self.fc(x))

model = Net().eval()

# 方式一：trace，需要示例输入
traced = torch.jit.trace(model, torch.randn(1, 8))
traced.save("net_traced.pt")

# 方式二：script，无需示例输入，保留控制流
scripted = torch.jit.script(model)
scripted.save("net_scripted.pt")

# 导出 ONNX
torch.onnx.export(
    model,
    torch.randn(1, 8),
    "net.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch"}, "output": {0: "batch"}},
    opset_version=17,
)

注意 model.eval() 这一步不可省略。训练态下的 Dropout、BatchNorm 行为与推理态不同，若忘记切换，导出的 ONNX 会带上训练语义，导致推理结果不稳定。

四、导出前准备清单与常见注意事项

检查项	正确做法	常见错误
模型状态	调用 `model.eval()`	忘记切换，Dropout/BatchNorm 语义错误
控制流	数据相关分支改用 `torch.where` 或 script	trace 只保留一条分支
输入形状	用 `dynamic_axes` 声明 batch 维	硬编码 batch=1，无法变长推理
算子支持	确认 opset 覆盖所用算子	自定义算子无 ONNX 映射
返回值	forward 只返回 Tensor 或 Tensor 元组	返回 dict / 自定义对象导致导出失败
设备	导出前统一在 CPU 上	CUDA 张量导致部分算子无法导出

建议在导出前先跑一次 torch.jit.trace 并用 traced.graph 打印图结构，肉眼确认算子序列与预期一致，再执行 ONNX 导出。这一步能提前暴露 80% 的分支丢失与算子缺失问题。

五、案例与复盘：一个典型的分支丢失事故

某检测模型在 forward 中写了 if self.training: ... else: ...，训练与推理走不同后处理。开发者直接 trace 后导出 ONNX，结果推理时后处理分支被固化成了训练分支，输出框数量异常。复盘结论是：凡是与输入数据或模型状态相关的分支，都不能依赖 trace 的默认行为，应改用 torch.jit.script，或在导出前把分支逻辑移出 forward，用外部配置固定。

另一个高频坑是 view 与 reshape 的混用。含 -1 的动态形状在 ONNX 中会生成 Reshape 且 shape 来自常量折叠，若上游维度是动态的，折叠会失败。导出前应尽量用 flatten 或显式 reshape 并配合 dynamic_axes 声明。

总结来说，导出前的准备工作可以归纳为四步：切 eval、清副作用、定形状、验算子。把 nn.Module 的 forward 整理成一张干净、可静态化的张量计算图，后续的 ONNX 导出与 MLIR 转换才会顺畅可控。