大模型算法实战教程

Appearance

09. PyTorch nn Module Basics | PyTorch nn.Module 基础 ​

难度： Medium | 环境： CPU-first | 标签： PyTorch, 模型封装, 状态管理 | 目标人群： Part 2-4 前置补课者

🚀 云端运行环境

本章节的实战代码可以点击以下链接在免费 GPU 算力平台上直接运行：

Open In ColabOpen In Studio (国内推荐：魔搭社区免费实例)

本页聚焦：会用 nn.Module 封装最小模型；会区分参数、缓冲区和普通属性；会读懂 state_dict() 和 parameters()。这一页开始把前面的 Tensor 和 Autograd 结果，正式收进一个模型对象里，让“能算”变成“能组织成模块”。

阅读顺序可以按这条线走：先看模型对象怎么成立，再看状态边界怎么划分，然后看 forward() 怎么串子模块，最后看 state_dict() 怎么保存和恢复。

关键词： nn.Module, Parameter, state_dict

前置阅读 ​

导语： 先看 0C 组页，把梯度边界和模型封装的边界对齐，再进入这一页会更顺。

相关阅读 ​

导语： 本页先把 nn.Module、参数注册和状态保存的最小判断讲清楚；如果想继续看 state_dict 的持久化细节，再顺着看下面这一页。

Q1：nn.Module 解决什么问题？ ​

模型不是一堆散落函数，而是一个有边界的对象。nn.Module 的作用，就是把参数注册、前向逻辑、子模块组合和状态保存统一起来。这里先把它当成 Part 2 里最常见的“模型对象入口”来看：你先确认这个对象怎么成立，再往下看它的状态和前向怎么组织。

python
import torch
import torch.nn as nn


class SimpleLinear(nn.Module):
    def __init__(self, in_features, out_features):
        super().__init__()
        # `Parameter` 会被自动注册进模型参数里。
        self.weight = nn.Parameter(torch.empty(out_features, in_features))
        self.bias = nn.Parameter(torch.zeros(out_features))
        nn.init.xavier_uniform_(self.weight)

    def forward(self, x):
        # `forward()` 定义计算逻辑，真正调用时写 `model(x)`。
        return x @ self.weight.t() + self.bias


class TwoLayerMLP(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        # 子模块会递归参与参数注册和状态管理。
        self.fc1 = SimpleLinear(input_dim, hidden_dim)
        self.act = nn.GELU()
        self.fc2 = SimpleLinear(hidden_dim, output_dim)

    def forward(self, x):
        # 典型的模块组合：线性层 -> 激活 -> 线性层。
        return self.fc2(self.act(self.fc1(x)))


def count_parameters(module):
    return sum(param.numel() for param in module.parameters())


model = TwoLayerMLP(4, 8, 2)
print('参数总量：', count_parameters(model))
print('named_parameters：', [name for name, _ in model.named_parameters()])

Q1验证：参数注册和前向输出是否正确？ ​

这里直接检查两件事：自定义层能不能跑，MLP 的输出 shape 和参数总量是否符合预期。

python
def test_simple_linear():
    layer = SimpleLinear(2, 1)
    with torch.no_grad():
        layer.weight.copy_(torch.tensor([[2.0, -1.0]]))
        layer.bias.copy_(torch.tensor([0.5]))
    x = torch.tensor([[3.0, 4.0]])
    y = layer(x)
    assert torch.allclose(y, torch.tensor([[2.5]]))
    print('✅ SimpleLinear 通过')


def test_two_layer_mlp():
    model = TwoLayerMLP(4, 8, 2)
    x = torch.randn(3, 4)
    y = model(x)
    assert y.shape == (3, 2)
    expected_params = (4 * 8 + 8) + (8 * 2 + 2)
    assert count_parameters(model) == expected_params
    print('✅ TwoLayerMLP 通过')


test_simple_linear()
test_two_layer_mlp()

Q2：什么时候必须区分 Parameter、buffer 和普通属性？ ​

当你开始关心模型里哪些东西会被训练、哪些东西只该保存、哪些东西只是临时配置时，就要先把状态边界看清。这里可以先按这个顺序判断：会不会训练更新 -> 会不会落进 state_dict() -> 只是临时配置还是长期状态。

python
class ToyModule(nn.Module):
    def __init__(self):
        super().__init__()
        self.weight = nn.Parameter(torch.tensor([1.0]))
        # `buffer` 常用于 mask、统计量、缓存这类“要保存但不训练”的状态。
        self.register_buffer('scale', torch.tensor([2.0]))
        self.name = 'toy'

    def forward(self, x):
        return x * self.weight * self.scale


m = ToyModule()
print('parameters：', [n for n, _ in m.named_parameters()])
print('buffers：', [n for n, _ in m.named_buffers()])
print('state_dict keys：', list(m.state_dict().keys()))

Q2验证：state_dict 是否同时包含参数和 buffer？ ​

这里确认 state_dict() 会包含参数和 buffer，但不会把普通属性当成可保存状态。

python
m = ToyModule()
keys = list(m.state_dict().keys())
assert 'weight' in keys
assert 'scale' in keys
assert 'name' not in keys
print('✅ state_dict 结构通过')

Q3：什么时候必须看 forward() 和子模块组合？ ​

当状态边界已经清楚后，再看 forward() 怎么把子模块串起来。这里真正重要的不是“写法”，而是“前向逻辑和模块边界是否清楚”。children()、named_children()、modules() 和 named_modules() 这几类接口，都是帮你快速看懂层级结构的入口。

python
x = torch.randn(2, 4)
model = TwoLayerMLP(4, 8, 2)
# 调用模块时写 `model(x)`，PyTorch 会帮你走到 `forward()`。
y = model(x)
print('输出 shape：', y.shape)
print('模块层级：', model)
print('named_children：', [name for name, _ in model.named_children()])

Q3验证：子模块组合是否正常？ ​

这里直接检查前向输出 shape 和模块结构，确认组合后的网络还能正常执行。

python
model = TwoLayerMLP(4, 8, 2)
x = torch.randn(5, 4)
y = model(x)
assert y.shape == (5, 2)
print('✅ forward 和子模块组合通过')

Q4：什么时候需要处理 state_dict 的保存和恢复？ ​

当模型对象和前向结构都看顺以后，最后就要把状态稳定落盘和恢复。这里记住：state_dict() 负责拿到当前模型状态，load_state_dict() 负责把状态恢复回去；它们关心的是权重和 buffer，不关心训练过程本身。

python
model = TwoLayerMLP(3, 6, 2)
# `state_dict()` 只保存状态，不保存训练过程。
before = {k: v.clone() for k, v in model.state_dict().items()}
buffer = model.state_dict()
model.load_state_dict(buffer)
after = model.state_dict()

for key in before:
    assert torch.allclose(before[key], after[key])
print('✅ state_dict roundtrip 通过')

Q4验证：保存和恢复是否不丢状态？ ​

这里只检查一件事：读出来再写回去，参数和 buffer 的值还在不在。你要把它理解成最小的“状态闭环”：只要状态没变，模型就能被稳定恢复。

python
model = TwoLayerMLP(3, 6, 2)
before = {k: v.clone() for k, v in model.state_dict().items()}
model.load_state_dict(model.state_dict())
after = model.state_dict()
for key in before:
    assert torch.allclose(before[key], after[key])
print('✅ state_dict 保持一致')

Released under the MIT License.

Copyright © 2024-present