nn.Module 架构设计与自定义
nn.Module 是 PyTorch 模型定义的基类，理解其内部机制对于构建复杂模型至关重要。
核心机制
python
class ResidualBlock(nn.Module):
def __init__(self, channels):
super().__init__()
self.conv1 = nn.Conv2d(channels, channels, 3, padding=1)
self.bn1 = nn.BatchNorm2d(channels)
self.conv2 = nn.Conv2d(channels, channels, 3, padding=1)
self.bn2 = nn.BatchNorm2d(channels)
def forward(self, x):
residual = x
out = F.relu(self.bn1(self.conv1(x)))
out = self.bn2(self.conv2(out))
return F.relu(out + residual)
# 参数管理
model = ResidualBlock(64)
print(sum(p.numel() for p in model.parameters())) # 总参数量
print(sum(p.numel() for p in model.parameters() if p.requires_grad)) # 可训练参数Module 的参数注册
只有通过 self.xxx = nn.Linear(...) 方式赋值的子模块才会被自动注册。直接用 Python 列表存储的模块不会被 parameters() 追踪，需使用 nn.ModuleList。