应继承 tf.keras.layers.Layer 并在 build() 中用 add_weight() 创建权重，避免在 __init__() 中手动调用 tf.Variable；call() 需声明 training 和 mask 参数以支持动态行为与掩码传播，并通过梯度检查验证变量注册与可微性。
如何继承 tf.keras.layers.Layer 实现自定义层
直接继承 tf.keras.layers.Layer 是最稳妥的方式，它能自动兼容 Keras 的构建逻辑（如权重初始化、build() 延迟调用、call() 签名规范）。别用 tf.keras.layers.Layer.__init__() 手动管理权重——那会绕过 Keras 的变量跟踪机制，导致 model.trainable_variables 漏掉你的参数。
实操建议：
- 在
__init__中只存超参（如units、activation），不创建权重 - 在
build(self, input_shape)中调用self.add_weight()创建可训练变量，Keras 会自动把它们加入trainable_weights -
call()必须接收inputs，可选接收training（用于 dropout/BatchNorm）、mask（用于序列掩码）
示例：一个带缩放的线性层
class ScaledDense(tf.keras.layers.Layer):
def __init__(self, units, scale_factor=1.0, **kwargs):
super().__init__(**kwargs)
self.units = units
self.scale_factor = scale_factor
<pre class="brush:php;toolbar:false;">def build(self, input_shape):
self.kernel = self.add_weight(
shape=(input_shape[-1], self.units),
initializer='glorot_uniform',
name='kernel'
)
self.bias = self.add_weight(
shape=(self.units,),
initializer='zeros',
name='bias'
)
def call(self, inputs, training=None):
return tf.matmul(inputs, self.kernel) * self.scale_factor + self.bias
立即学习“Python免费学习笔记（深入）”；
为什么不能在 __init__ 里直接调用 tf.Variable
手动用 tf.Variable 创建权重会导致 Keras 无法识别该变量属于哪一层，进而引发三个问题：保存模型时漏掉该变量、model.summary() 不显示、model.trainable_variables 不包含它。更隐蔽的是，当层被重复使用（如在多个分支中调用同一实例），tf.Variable 会复用而非复制，造成意外的参数共享。
正确做法始终是用 self.add_weight()，它内部调用 tf.Variable 但额外注册了所有权关系。
常见错误现象：
-
model.save('m.h5')后加载发现权重全为初始值 -
len(model.trainable_variables)比预期少 - 多 GPU 分布式训练时报错
ValueError: Variable ... is not in the list of trainable variables
如何支持 training=True/False 和 mask 传播
如果你的层有随机行为（如 dropout）或状态依赖（如 BatchNorm），必须显式处理 training 参数；如果输入可能带 mask（如 LSTM 输出），需检查并传递 mask。Keras 不会自动帮你转发这两个参数——除非你在 call() 签名里明确声明，并在内部逻辑中响应它们。
实操要点：
- 不要写
def call(self, inputs):，而要写def call(self, inputs, training=None, mask=None): - 对 dropout 类操作，用
tf.nn.dropout(x, rate, training=training)，而不是硬编码training=True - 若需要向下传递 mask（如自定义 RNN 单元），返回时加
mask=mask到输出元组，或调用self.compute_mask(inputs, mask)（需重写该方法）
性能影响：忽略 training 会导致推理时仍执行随机丢弃，严重拖慢速度且结果不可复现。
如何验证自定义层是否被正确纳入计算图和梯度流
最简单的验证方式不是跑通训练，而是检查三件事：变量是否注册、前向是否可微、反向是否连通。别等到训练 loss 不下降才排查。
快速验证步骤：
- 构造一个最小输入
x = tf.random.normal((2, 10))，调用layer(x)，确认无报错 - 检查
layer.trainable_weights是否非空且形状符合预期 - 用
with tf.GradientTape() as tape:包裹y = layer(x)，再调用tape.gradient(y, layer.trainable_weights)，确保返回的梯度列表不全为None
容易被忽略的点：如果层内部用了未被 tf.function 装饰的 Python 控制流（如 if x > 0:），梯度可能中断；此时应改用 tf.cond 或确保所有分支都参与计算图构建。
