Python 现代开发学习笔记

(十二)nn.Module 与参数登记

前面的 w 与 b 是两个独立张量。参数多起来后,需要一起移动设备、清理梯度、保存和恢复,可以把它们放进 nn.Module。

# 项目文件:demos/module_parameters.py
import torch
from torch import nn


class TinyLinear(nn.Module):
    def __init__(self):
        super().__init__()
        self.weight = nn.Parameter(torch.tensor([[2.0], [3.0]]))
        self.bias = nn.Parameter(torch.tensor([5.0]))
        self.register_buffer("mean", torch.zeros(2))  # 随模型保存,但不参加优化。
        self.description = "两个输入特征"  # 普通属性不是训练参数。

    def forward(self, x):
        return (x - self.mean) @ self.weight + self.bias


model = TinyLinear()
print("预测", model(torch.tensor([[10.0, 2.0]])).item())
print("参数名", [name for name, _ in model.named_parameters()])
print("状态名", list(model.state_dict()))

先调用 super().__init__() 准备模型的登记机制。赋给实例属性的 nn.Parameter 会作为可训练参数登记,model.parameters() 可以交给优化器,named_parameters() 同时给出名称。

forward 写本次前向计算;平时调用 model(x),由 Module 的调用流程进入 forward。例中仍计算 10×2+2×3+5=31,只是参数现在由一个对象共同管理。

mean 通过 register_buffer 登记为缓冲数据,会随模型移动和保存,但不会进入优化器的参数列表。description 则是普通 Python 属性。因此参数名只有 weight、bias,state_dict 中还会出现 mean。

模型把 weight 和 bias 登记为参数,把 mean 登记为缓冲区,普通说明属性不进入权重状态,调用模型再进入 forward

完整回归模型直接使用现成的线性层:

# 项目文件:training_project/src/readingtorch/models.py
from torch import nn


class ReadingRegressor(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(2, 1)  # 两个输入特征,输出一个标准化页数。

    def forward(self, x):
        return self.linear(x)

nn.Linear(2, 1) 保存两项权重和一项偏置,把输入的最后一维从二变为一。PyTorch 内部权重按 (输出特征数, 输入特征数) 保存,这里是 (1, 2);前向等价于输入乘权重的转置,再加偏置。

笔记正文与全部题解 · Ctrl / ⌘ K