2.4 自定义层与模块：自己造货架

2.4 自定义层与模块：自己造货架

本节摘要：预置层不够用的那天，你要能自己造一个。本节把 2.2 节的登记机制反过来用：只要参数用 nn.Parameter 声明、子层作为属性挂载，你自己写的类就与框架层完全同权——能进 Sequential、能被优化器接管、能随 state_dict 保存。

框架的边界在哪里

组建营的最后一课。什么时候需要自定义？三种典型情形：论文里的新结构框架尚未收录；你要给某个标准结构加一点私货（比如自定义初始化、特殊缩放）；你想把一组固定搭配的层打包成一个语义单元，让主模型代码更干净。三种情形的解法是同一个——写一个继承 nn.Module 的类。

自定义层的自由度来自一个事实：forward 里你可以写任意 Python 代码。条件分支、循环、调用其他库，都合法。这是动态图框架的底气（导读里提过），静态图框架里这类写法要小心处理。

图 2-2：自定义层与框架设施的兼容路径

造一个能用的自定义层

以一个实战里真实出现过的需求为例：给特征做"可学习的逐通道缩放与平移"。它比 BatchNorm 简单——不统计批分布，只靠两个可学习向量逐通道仿射变换，序列任务里偶有妙用。

import torch
import torch.nn as nn

class ChannelScale(nn.Module):
    """逐通道可学习缩放与平移：输入 形状 x 通道数"""
    def __init__(self, channels):
        super().__init__()
        self.gamma = nn.Parameter(torch.ones(channels))    # 初始为1：起始等价于恒等变换
        self.beta = nn.Parameter(torch.zeros(channels))    # 初始为0：不额外平移

    def forward(self, x):
        # x 形状 batch x channels x ...；把参数 reshape 成可广播的形状
        shape = [1, -1] + [1] * (x.dim() - 2)
        return x * self.gamma.view(shape) + self.beta.view(shape)

layer = ChannelScale(8)
x = torch.randn(4, 8, 3, 3)
y = layer(x)
print("输出形状:", tuple(y.shape))
print("入册参数:", [(n, tuple(p.shape)) for n, p in layer.named_parameters()])
# 恒等性检查：初始 gamma=1、beta=0 时应当完全透传
print("初始为恒等变换:", torch.allclose(y, x))
​

输出：

输出形状: torch.Size([4, 8, 3, 3])
入册参数: [('gamma', (8,)), ('beta', (8,))]
初始为恒等变换: True
​

三个工程细节都藏在这二十行里。初始值选"恒等"（gamma 全 1、beta 全 0），让自定义层插入现有网络时一开始不改变任何行为，训练再慢慢学出偏移——这是加自定义层的稳妥姿势，远好于随机初始化把已经训练好的网络当场搅乱。view(shape) 那一行用 1.3 节的广播知识把参数变形到可广播，是自定义层最常见的样板代码。

完整案例：打包一个"卷积块"

背景：6.1 节的视觉模型里，"卷积加归一化加激活"这组三件套出现了两次，主模型代码又长又容易抄错其中一行的顺序。

操作：把三件套打包成一个自定义块，让主模型只看得到语义单元。

import torch
import torch.nn as nn

class ConvBlock(nn.Module):
    """卷积 + BatchNorm + ReLU 的标准三件套"""
    def __init__(self, in_ch, out_ch):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(in_ch, out_ch, 3, padding=1),
            nn.BatchNorm2d(out_ch),
            nn.ReLU(),
        )
    def forward(self, x):
        return self.net(x)

model = nn.Sequential(ConvBlock(1, 8), nn.MaxPool2d(2), ConvBlock(8, 16), nn.MaxPool2d(2))
x = torch.randn(4, 1, 8, 8)
print("输出形状:", tuple(model(x).shape))
print("顶层子模块:", [n for n, _ in model.named_children()])
​

输出：

输出形状: torch.Size([4, 16, 2, 2])
顶层子模块: ['0', '1', '2', '3']
​

结果：主模型从六行缩到一行语义，而嵌套的参数照常被登记——named_parameters() 依然能枚举全部参数，因为登记表是递归的（子模块的 _modules 里再挂子模块）。

解读：自定义的价值一半在能力扩展，一半在语义打包。当 reviewers 读你的模型代码时，看到 ConvBlock 比看到三行裸层信息密度高得多。另外注意形状账：两次池化把 8×8 压到 2×2，通道从 1 涨到 16，这种"空间换通道"是视觉网络的典型轮廓。

变式：给 ConvBlock 加一个可选的残差捷径（输入输出同形状时把输入直接加到输出上），只需要在 forward 里多一个 if 分支加一行 return x + self.net(x)。动手改一次，你会对"forward 里任意 Python"这句话有实感。

本节要点回顾

• 自定义层三件事：继承 nn.Module 必做、参数用 nn.Parameter 必做、常量用 buffer 按需；
• 初始值设计成恒等变换，让新层插入时不搅乱既有行为；
• 登记表递归生效：嵌套的子模块参数照样能被枚举、优化、保存；
• forward 里任意 Python 合法，这是动态图框架给自定义结构的空间。

下一节进入第 3 章：把数据装进流水线，让远征队开拔。