3、PyTorch模型基础：构建简单神经网络模型，理解torch.nn.Module

3.1 核心原理：为什么一切模型都从nn.Module开始

在PyTorch中，torch.nn.Module是所有神经网络模块的基类。无论是线性层、卷积层、Transformer编码器，还是你即将导出的完整ONNX模型，最终都继承自这个类。nn.Module的核心价值在于它提供了一套参数管理与前向传播的标准化框架。当你定义一个子类并实现forward()方法时，PyTorch会自动完成以下关键工作：

  • 参数注册：通过self.xxx = nn.Parameter(...)或直接使用nn.Linear等层时，参数会被自动注册到模块的parameters()迭代器中，供优化器更新。
  • 子模块追踪：嵌套的nn.Module会被记录在children()和modules()中，使得整个模型的结构可被遍历和序列化。
  • 训练/评估模式切换：调用model.train()或model.eval()会递归影响所有子模块（如Dropout、BatchNorm的行为差异）。
  • 设备迁移：model.to(device)可将所有参数和缓冲区一次性移动到GPU或CPU。

理解nn.Module的机制，是后续进行ONNX导出的前提——因为ONNX导出过程本质上是追踪forward()执行时的数据流图，而非静态解析Python代码。因此，forward()中使用的操作必须是PyTorch可追踪的（如张量运算、nn.functional函数），而不能包含Python原生控制流（如if依赖张量值、for循环依赖张量形状）。

关键认知：ONNX导出时，PyTorch会“执行”一次你的forward()，记录所有张量操作形成计算图。这意味着forward()中的Python逻辑（如print、assert）不会出现在ONNX中，但会影响导出时的执行路径。
nn.Module 内部机制与ONNX导出关系 nn.Module 基类 __init__() 中定义层 self.fc1 = nn.Linear(784, 128) parameters() 迭代器 weight: [128, 784], bias: [128] forward() 定义计算 x = torch.relu(self.fc1(x)) train() / eval() 影响Dropout、BatchNorm行为 ONNX 导出追踪 torch.onnx.export() model, dummy_input, ... 执行 forward() 记录张量操作序列 生成计算图 节点 = 算子，边 = 张量 输出 model.onnx ONNX 文件内容 graph 节点列表 Gemm, Relu, Softmax... initializer 初始权重 fc1.weight, fc1.bias graph.input / output input: [1,784], output: [1,10] opset 算子版本 opset_version=17 继承 生成 注意：ONNX 不包含Python控制流，只包含张量运算图。若forward中有if/for依赖张量值，需用torch.onnx.is_onnx_export()处理

3.2 方法流程：构建一个用于MNIST分类的简单神经网络

下面我们构建一个包含两个全连接层和一个ReLU激活函数的简单神经网络。该模型将作为后续ONNX导出的基础示例。整个流程分为三步：定义模型结构、实例化并检查参数、执行一次前向传播验证。

第一步：定义模型类

import torch
import torch.nn as nn
import torch.nn.functional as F

class SimpleMLP(nn.Module):
    """一个简单的多层感知机，用于MNIST分类（28x28=784输入，10类输出）"""
    def __init__(self):
        super(SimpleMLP, self).__init__()
        # 定义网络层
        self.fc1 = nn.Linear(784, 128)   # 输入层 -> 隐藏层
        self.fc2 = nn.Linear(128, 64)    # 隐藏层 -> 第二个隐藏层
        self.fc3 = nn.Linear(64, 10)     # 隐藏层 -> 输出层
        # Dropout用于防止过拟合（训练时随机失活）
        self.dropout = nn.Dropout(0.2)

    def forward(self, x):
        # 前向传播逻辑
        x = torch.flatten(x, start_dim=1)  # 将 [batch, 1, 28, 28] 展平为 [batch, 784]
        x = F.relu(self.fc1(x))
        x = self.dropout(x)                # 训练时生效，eval()时自动关闭
        x = F.relu(self.fc2(x))
        x = self.fc3(x)                    # 输出层不加激活，留给损失函数处理
        return x

第二步：实例化并检查模型结构

model = SimpleMLP()
print(model)

# 输出：
# SimpleMLP(
#   (fc1): Linear(in_features=784, out_features=128, bias=True)
#   (fc2): Linear(in_features=128, out_features=64, bias=True)
#   (fc3): Linear(in_features=64, out_features=10, bias=True)
#   (dropout): Dropout(p=0.2, inplace=False)
# )

# 统计参数数量
total_params = sum(p.numel() for p in model.parameters())
trainable_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f"总参数量: {total_params:,} | 可训练参数量: {trainable_params:,}")
# 输出：总参数量: 109,386 | 可训练参数量: 109,386

第三步：执行前向传播验证

# 创建模拟输入（MNIST图像：1通道，28x28）
dummy_input = torch.randn(1, 1, 28, 28)

# 切换到评估模式（关闭Dropout）
model.eval()
with torch.no_grad():
    output = model(dummy_input)

print(f"输入形状: {dummy_input.shape}")
print(f"输出形状: {output.shape}")
print(f"输出张量（前5个值）: {output[0, :5]}")

# 输出：
# 输入形状: torch.Size([1, 1, 28, 28])
# 输出形状: torch.Size([1, 10])
# 输出张量（前5个值）: tensor([-0.0231,  0.0452,  0.0128, -0.0310,  0.0189])
重要提醒：在导出ONNX之前，务必调用model.eval()。否则，Dropout层和BatchNorm层会保留训练行为，导致导出的模型推理结果不稳定。ONNX模型通常用于推理部署，因此必须确保其处于评估模式。

3.3 案例场景：理解nn.Module的三种典型用法

在实际项目中，nn.Module的使用方式直接影响模型的可维护性和ONNX导出的顺利程度。以下是三种常见场景的对比：

场景 代码模式 ONNX导出影响
顺序容器
（快速堆叠层）		 
self.layers = nn.Sequential(
    nn.Linear(784, 128),
    nn.ReLU(),
    nn.Linear(128, 10)
)
 导出最稳定，结构清晰，推荐用于简单模型。
自定义forward
（灵活控制数据流）		 
def forward(self, x):
    x = self.fc1(x)
    x = F.relu(x)
    x = self.dropout(x)
    return self.fc2(x)
 可导出，但需确保所有操作都是可追踪的（不使用Python原生控制流）。
多输入/多输出
（复杂模型）		 
def forward(self, x1, x2):
    h = torch.cat([x1, x2], dim=1)
    return self.fc(h)
 导出时需要指定多个dummy_input，且ONNX中需定义多个输入节点。

对于本课程后续的ONNX导出实战，我们推荐使用自定义forward模式，因为它最接近真实项目中的模型写法，同时也能清晰展示数据流。但要注意：forward()中应避免使用print、if x.sum() > 0这类依赖张量值的Python控制流，否则导出时会报错或产生错误图。

3.4 误区与复盘：nn.Module使用中的常见陷阱

在构建模型和准备ONNX导出时，以下四个误区最容易导致问题。我们逐一分析并给出解决方案。

误区一：忘记调用super().__init__()
如果子类的__init__()中没有调用super().__init__()，则nn.Module的初始化逻辑不会执行，导致参数无法注册。此时调用model.parameters()会返回空迭代器，训练时优化器无参数可更新。

误区二：在forward()中使用Python原生控制流

# 错误示例：依赖张量值的if语句
def forward(self, x):
    if x.mean() > 0:   # 导出ONNX时会报错
        x = self.fc1(x)
    else:
        x = self.fc2(x)
    return x

# 正确做法：使用torch.where或保持静态结构
def forward(self, x):
    # 使用torch.where实现条件选择（可追踪）
    x = torch.where(x.mean() > 0, self.fc1(x), self.fc2(x))
    return x

误区三：混淆nn.Module与nn.functional

  • nn.Linear是一个模块，内部包含可学习的权重参数，需要注册到模型中。
  • F.relu是一个函数，无参数，直接调用即可。
  • 如果误将F.linear当作模块使用，会导致参数无法被优化器更新。正确做法是：有参数的层用nn.xxx，无参数的激活函数用F.xxx。

误区四：导出前未切换eval()模式

以下代码展示了错误与正确的对比：

# 错误：导出前未调用eval()
torch.onnx.export(model, dummy_input, "model.onnx")  # Dropout仍在训练模式

# 正确：先eval()再导出
model.eval()
torch.onnx.export(model, dummy_input, "model.onnx",
                  input_names=["input"], output_names=["output"],
                  opset_version=17)

复盘总结：构建nn.Module模型时，请始终遵循以下清单：

  1. 在__init__()中先调用super().__init__()。
  2. 所有含参数的层（Linear、Conv2d等）用self.xxx = nn.xxx()定义。
  3. 无参数的激活函数（ReLU、Sigmoid）在forward()中用F.xxx()调用。
  4. 导出ONNX前，必须调用model.eval()并配合torch.no_grad()。
  5. 使用dummy_input验证前向传播输出形状是否符合预期。
← 上一章 📖 返回目录 下一章 →