3、自定义层基础：继承 tf.keras.layers.Layer、__init__ 与 call 方法、build 方法的作用
好，咱们今天聊聊自定义层。说实话，这是 TensorFlow 进阶路上绕不开的一道坎。很多同学学完 Sequential 和 Functional API 后，觉得模型搭建不过如此。但一旦遇到稍微复杂点的需求——比如共享权重、动态结构、或者需要精细控制前向传播——就会发现内置层不够用了。
这时候，你就得自己动手写层了。
3.1 为什么要自定义层？
我刚开始用 TensorFlow 时，也天真地以为内置层能搞定一切。直到有一次做 NLP 项目，需要实现一个带掩码的注意力层，翻遍官方文档都没找到现成的。嗯，那时候我才意识到：框架给你的是积木，但真正的创造力在于你自己搭积木。
自定义层的核心场景包括：
- 特殊运算：比如自定义激活函数、归一化方式
- 权重共享：多个位置共用同一组可训练参数
- 动态结构：根据输入形状决定内部结构
- 调试与监控：在 forward 过程中插入打印或统计
一句话总结：当你发现「搭积木」搭不出想要的效果时，就该自己造积木了。
3.2 继承 tf.keras.layers.Layer
所有自定义层的起点，都是继承 tf.keras.layers.Layer。这个基类帮你处理了参数管理、设备分配、序列化等脏活累活。你只需要关注三个核心方法：__init__、build、call。
先看一个最简单的例子：
import tensorflow as tf
class MyDense(tf.keras.layers.Layer):
def __init__(self, units, activation=None):
super(MyDense, self).__init__()
self.units = units
self.activation = tf.keras.activations.get(activation)
def build(self, input_shape):
self.w = self.add_weight(
shape=(input_shape[-1], self.units),
initializer='random_normal',
trainable=True,
name='kernel'
)
self.b = self.add_weight(
shape=(self.units,),
initializer='zeros',
trainable=True,
name='bias'
)
def call(self, inputs):
z = tf.matmul(inputs, self.w) + self.b
return self.activation(z)这个层实现了全连接层的基本功能。你可能会说：「这不就是 Dense 层吗？」对，但重点是——你完全掌控了内部逻辑。
3.3 __init__ 方法：初始化配置
__init__ 负责接收用户传入的超参数。我个人习惯在这里只做两件事：
- 保存用户传入的参数（如 units、activation）
- 调用父类的
__init__
注意：不要在 __init__ 里创建权重。为什么？因为此时你还不知道输入的形状。你想想看，如果用户传入的 input_shape 是 (None, 64)，你根本不知道第二个维度是多少，怎么创建权重矩阵？
我的经验：__init__ 里只存「配置」，不存「数据」。这样层才能灵活适应不同的输入形状。
3.4 build 方法：延迟创建权重
build 方法才是真正创建权重的地方。它会在第一次调用 call 之前被自动执行，并且会传入输入的形状信息。
为什么要有这个延迟机制？说白了，就是解耦配置与形状。你可以在不知道输入维度的情况下先定义层，等实际数据来了再自动适配。
我在项目中遇到过这样一个坑：早期版本里，我习惯在 __init__ 里直接创建权重，结果换了个数据集，输入维度变了，整个模型都得重写。后来改用 build 方法，再也没出过这种问题。
build 方法的核心职责：
- 使用
self.add_weight()创建可训练参数 - 根据
input_shape动态计算参数维度 - 设置
self.built = True（基类自动处理）
注意：如果你在 build 里创建了多个权重，记得给每个权重起个有意义的 name。否则在模型保存和加载时，你可能会看到一堆莫名其妙的变量名。
3.5 call 方法：前向传播逻辑
call 是层的核心——它定义了输入如何变成输出。这里你可以做任何 TensorFlow 支持的操作：矩阵乘法、卷积、注意力机制、甚至调用其他层。
call 方法的几个要点：
- 接收
inputs作为第一个参数 - 可以接收额外的参数（如 training 标志）
- 返回处理后的张量
举个例子，假设我们要实现一个带 dropout 的自定义层：
class MyDropoutLayer(tf.keras.layers.Layer):
def __init__(self, rate):
super(MyDropoutLayer, self).__init__()
self.rate = rate
def call(self, inputs, training=None):
if training:
return tf.nn.dropout(inputs, rate=self.rate)
return inputs看到了吗？training 参数让这个层在训练和推理时表现不同。这种灵活性是内置层无法直接提供的。
3.6 三个方法的协作流程
咱们捋一下整个生命周期：
- 实例化：调用
__init__，保存配置 - 第一次调用：框架检测到
self.built == False，自动调用build(input_shape) - 前向传播：执行
call(inputs)，返回结果 - 后续调用：跳过 build，直接执行 call
这个流程设计得非常巧妙。你想想看，如果每次调用都重新 build，那效率得多低？但如果不在 build 里创建权重，又没法处理动态形状。所以 TensorFlow 选择了「一次 build，多次 call」的策略。
核心思想：__init__ 负责「是什么」，build 负责「长什么样」，call 负责「做什么」。
3.7 实战：一个带正则化的自定义层
最后，咱们结合前面讲的内容，写一个稍微复杂点的层。这个层会在全连接的基础上，加上 L2 正则化：
class RegularizedDense(tf.keras.layers.Layer):
def __init__(self, units, l2_lambda=0.01):
super(RegularizedDense, self).__init__()
self.units = units
self.l2_lambda = l2_lambda
def build(self, input_shape):
self.w = self.add_weight(
shape=(input_shape[-1], self.units),
initializer='glorot_uniform',
regularizer=tf.keras.regularizers.L2(self.l2_lambda),
trainable=True,
name='kernel'
)
self.b = self.add_weight(
shape=(self.units,),
initializer='zeros',
trainable=True,
name='bias'
)
def call(self, inputs):
return tf.matmul(inputs, self.w) + self.b注意看，我在 add_weight 里传入了 regularizer 参数。这样，当这个层被加入模型后，正则化损失会自动被 model.losses 收集。我曾经踩过这个坑——手动计算正则化损失，结果忘了加进去，模型训练出来过拟合得一塌糊涂。
3.8 小结
自定义层其实没那么神秘。你只需要记住三个方法的分工：
| 方法 | 职责 | 执行时机 |
|---|---|---|
| __init__ | 保存超参数配置 | 实例化时 |
| build | 创建可训练权重 | 第一次调用前 |
| call | 定义前向传播逻辑 | 每次调用时 |
嗯，掌握了这个基础，后面咱们就可以玩更花哨的了——比如自定义训练循环、梯度裁剪、甚至自定义优化器。但万丈高楼平地起，先把层写好，后面的路就好走了。
下一章，咱们聊聊如何用自定义层搭建完整的模型，以及如何在训练过程中监控每一层的输出。到时候我会分享一个我踩过的坑——关于梯度消失的调试经验，保证让你少走弯路。