(十二)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。
完整回归模型直接使用现成的线性层:
# 项目文件: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);前向等价于输入乘权重的转置,再加偏置。