nn.Module：构建神经网络的积木

从一块乐高底板说起

神经网络由很多层组成：线性层、卷积层、激活函数……一层套一层。如果每层都靠你手动管参数、管梯度，模型一大就乱套。手动管理的痛苦，前几章你应该已经体会过了：更新权重要挨个写，清梯度要挨个写，稍不留神就漏一个。

PyTorch 想了个办法：提供一个统一的底座，叫 nn.Module。所有层、所有模型都建在这个底座上。它负责三件事：登记参数、记录子层、附赠常用工具。

打个比方：nn.Module 像乐高底板，nn.Linear、nn.Conv2d 这些层是积木块。你把积木块插在底板上，底板自动知道「我身上有哪些块」。搬家、清点、保养，它都能代办。

一个最小的网络类

先看代码，再逐行解释：

import torch
import torch.nn as nn

class TinyNet(nn.Module):
    def __init__(self):
        super().__init__()          # 初始化底座，千万别漏
        self.fc1 = nn.Linear(4, 8)  # 第一层：4 进 8 出
        self.fc2 = nn.Linear(8, 2)  # 第二层：8 进 2 出

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = self.fc2(x)
        return x

model = TinyNet()
print(model)

输出：

TinyNet(
  (fc1): Linear(in_features=4, out_features=8, bias=True)
  (fc2): Linear(in_features=8, out_features=2, bias=True)
)

两个方法各管一半：

1. __init__ 负责「摆积木」，定义网络里有哪些层。
2. forward 负责「数据怎么流」，规定数据按什么顺序穿过这些层。

调用时不用写 model.forward(x)，直接 model(x) 就行。模型对象像函数一样可调用，PyTorch 在背后帮你转交给 forward。这里有个坑：第一行 super().__init__() 漏了，底座没初始化，参数就登不了记。这个错很隐蔽，因为代码照样能跑，只是后面训练死活不收敛。

打印模型还有额外好处：结构一目了然。Linear(in_features=4, out_features=8, bias=True) 直接告诉你每层的进出维度和是否有偏置。排查形状问题的时候，先 print(model) 往往比看半天代码快。

为什么参数能自动登记

普通张量放进模型，PyTorch 不认；nn.Linear 这种层自带权重（weight）和偏置（bias），赋值给 self.fc1 的那一刻，底座就把它的参数纳入管理了。

想验证？一行代码：

print(len(list(model.parameters())))  # 4：两层各有权重和偏置

parameters() 返回模型里所有可训练参数。后面章节你会看到，优化器就是靠它拿到参数清单，反向传播后统一更新的。同理，model.zero_grad() 一键清空所有参数的梯度，不用你挨个写清零语句。

模块还能套模块

nn.Module 的登记是递归的：你的模型里挂了一个子模块，子模块里又挂了孙模块，底座会一层层查下去，把最底层的参数全部收进 parameters()。这意味着你可以把网络拆成小块，各自封装，再拼成整体：

class Block(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Linear(8, 8)

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

class BigNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.block1 = Block()
        self.block2 = Block()
        self.out = nn.Linear(8, 2)

    def forward(self, x):
        return self.out(self.block2(self.block1(x)))

这种「小积木拼大积木」的写法，是 PyTorch 项目最常见的组织方式。每个 Block 自己管好自己的参数，父模块什么都不用操心。这也是 nn.Module 名字里 Module 的含义：它本身就是个可嵌套的模块。

底座附赠的实用工具

除了 parameters() 和 zero_grad()，还有两个高频工具：

1. model.to(device)：把整个模型（所有参数）一次性搬到指定设备。第 8 章讲过设备管理，model.to("cuda") 就是模型级搬家，不用挨个层搬。
2. model.train() / model.eval()：切换训练和评估状态。像 nn.Dropout、nn.BatchNorm 这类层，两种状态下行为不同。评估前别忘了切 eval()。

model.to("cuda")          # 搬到 GPU
model.train()             # 进入训练模式
# ... 训练若干轮 ...
model.eval()              # 进入评估模式

自定义参数：nn.Parameter

大多数情况用现成层就够了。真要自己造参数，就包一层 nn.Parameter：

class CustomNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.w = nn.Parameter(torch.randn(4, 2))
        self.b = nn.Parameter(torch.zeros(2))

    def forward(self, x):
        return x @ self.w + self.b

nn.Parameter 本质是张量，但自带「我是要训练的」标记。省掉它直接算 x @ self.w 也能前向传播，但那时 self.w 只是普通属性，parameters() 里看不到，训练时不会被更新。我早期就栽过这个跟头：模型明明能跑，loss 就是不降，查半天才发现参数没登记。

nn 层与 functional 函数

你可能见过两种写法：nn.ReLU() 和 F.relu(x)，效果一样，区别在「有没有状态」。

• nn.ReLU() 是类，实例化后可以当积木放进容器。
• F.relu() 是纯函数，直接对张量操作，什么都不保存。

ReLU 没参数，两种写法都行。但 nn.Linear 这种带权重的层，就别用 F.linear 手动传权重了，交给层对象管理省心得多。记住一个原则：有参数的用层（nn.*），没参数的操作两种随意。

小结

• nn.Module 是神经网络的统一底座，负责登记参数、管理子层。
• __init__ 摆层，forward 定流向，super().__init__() 别忘。
• 模型可直接调用：model(x)，打印它还能看到完整结构。
• 模块可嵌套，小块拼大块是标准组织方式。
• parameters()、zero_grad()、to(device)、train()/eval() 都是底座送的。
• 自定义参数用 nn.Parameter 包一下，才会被登记。

下一章我们认识常用的网络层和容器，看看怎么把积木拼得更快。