子类化Keras Layer却触发Lambda层变量未追踪警告的疑问

首先得明确核心问题：你当前继承的是tf.Module，而非Keras官方推荐的tf.keras.layers.Layer——这就是警告出现的根源。

警告背后的逻辑 #

Keras的Model在处理组件时，会优先识别tf.keras.layers.Layer的子类。当你传入一个tf.Module实例时，Keras并没有把它当成标准的Keras层，而是将其包装成Lambda层来处理。

Lambda层的定位是处理简单的无状态计算，它本身不会自动跟踪内部的可训练变量。所以当Keras检测到这个Lambda层里用到了self.w和self.b这两个可训练变量，但这些变量没有被Keras的变量跟踪系统捕获时，就会抛出警告：它在提醒你，这种写法可能导致变量无法被正确管理——比如训练时不更新、模型保存后加载丢失变量、分布式训练出现异常等。

怎么解决这个问题？ #

把你的类改成继承tf.keras.layers.Layer，并且遵循Keras层的规范实现：

1. 继承tf.keras.layers.Layer（它本身已经继承了tf.Module，完全保留Module功能的同时，还多了Keras专属特性）
2. 用self.add_weight()创建可训练变量（Keras会自动跟踪这些变量）
3. 重写call()方法代替__call__（Keras的Layer已经封装了__call__的逻辑，包括training参数处理、输入形状验证等）

修改后的代码如下：

import tensorflow as tf
class MyCustomLayer(tf.keras.layers.Layer):
    def __init__(self, filters, kernel_size, padding, strides, activation, kernel_initializer, bias_initializer, use_bias):
        super(MyCustomLayer, self).__init__()
        self.filters = filters
        self.kernel_size = kernel_size
        self.activation = activation
        self.padding = padding
        self.kernel_initializer = kernel_initializer
        self.bias_initializer = bias_initializer
        self.strides = strides
        self.use_bias = use_bias
        # 无需手动传入in_channels，后续build方法会根据输入形状自动推断

    def build(self, input_shape):
        # 从输入形状中获取输入通道数
        in_channels = input_shape[-1]
        # 用add_weight创建可训练变量，自动被Keras跟踪
        self.w = self.add_weight(
            shape=(*self.kernel_size, in_channels, self.filters),
            initializer=self.kernel_initializer,
            trainable=True,
            name="conv_kernel"
        )
        if self.use_bias:
            self.b = self.add_weight(
                shape=(self.filters,),
                initializer=self.bias_initializer,
                trainable=True,
                name="conv_bias"
            )
        super().build(input_shape)

    def call(self, inputs, training=None):
        x = tf.nn.conv2d(inputs, filters=self.w, strides=self.strides, padding=self.padding)
        if self.use_bias:
            x = tf.add(x, self.b)
        x = self.activation(x)
        return x

x = tf.keras.Input(shape=(28, 28, 3))
# 这里不需要传入in_channels了，build方法会自动获取
y = MyCustomLayer(
    filters=16, kernel_size=(3, 3), strides=(1, 1), activation=tf.nn.relu,
    padding='VALID', kernel_initializer=tf.initializers.GlorotUniform(),
    bias_initializer=tf.initializers.Zeros(), use_bias=True)(x)
model = tf.keras.Model(inputs=x, outputs=y)

额外优化点 #

我去掉了__init__里的in_channels参数，改用build方法自动从输入形状中推断——这是Keras层的标准做法，让你的层更灵活，无需手动指定输入通道数。

总结警告的意义 #

这个警告不是说你的代码完全不能运行，而是在提醒你：当前写法不符合Keras的层规范，可能带来变量管理上的隐患。改用tf.keras.layers.Layer子类化，就能让Keras正确识别和管理你的层的变量，消除警告，同时获得Keras层的所有特性（比如模型保存/加载、输入形状验证、分布式训练支持等）。

内容的提问来源于stack exchange，提问作者Nicolas Gervais