如何在自定义Keras层的call方法中创建不可训练的tf.Variable

你的代码存在两个核心问题导致梯度报错：

1. 在@tf.function修饰的call方法内动态创建tf.Variable，这违反了TensorFlow图模式的规则——变量必须在图构建阶段（__init__或build方法）创建，不能在图运行时动态生成。
2. 返回的self.outs是trainable=False的变量，梯度计算时会将其视为常数，导致梯度无法传导到训练变量self.w，最终引发"No gradients provided for any variable"错误。

方案1：移除不必要的状态变量（最简修复） #

如果你的业务逻辑只是每次返回x + self.w，不需要维护self.outs的状态，直接简化代码即可：

def __init__(self, filters, conv_size, lag, **kwargs):
    super(Upscaler, self).__init__(**kwargs)
    self.w = self.add_weight(
        shape=(1, 128, 128, 3),
        trainable=True,
        name='conv_weight'
    )

@tf.function
def call(self, x):
    return x + self.w

这样梯度能正常传导到self.w，完全规避状态变量带来的问题。

方案2：正确维护状态变量（保留状态需求） #

如果确实需要维护self.outs的状态（比如累积计算），需按以下规则修改：

• 在build方法中创建非训练状态变量（输入形状在__init__阶段未知）
• 返回张量而非状态变量本身，确保梯度路径不被截断

修改后的代码：

def __init__(self, filters, conv_size, lag, **kwargs):
    super(Upscaler, self).__init__(**kwargs)
    self.w = self.add_weight(
        shape=(1, 128, 128, 3),
        trainable=True,
        name='conv_weight'
    )
    self.outs = None
    # 用tf.Variable存储初始化状态，避免图模式下Python变量被固化
    self._initialized = tf.Variable(False, trainable=False, name='initialized')

def build(self, input_shape):
    # 根据输入形状创建非训练状态变量
    self.outs = self.add_weight(
        shape=input_shape[1:],  # 忽略batch维度
        initializer='zeros',
        trainable=False,
        name='state_outs'
    )
    super().build(input_shape)

@tf.function
def call(self, x):
    # 第一次调用初始化状态
    if tf.logical_not(self._initialized):
        self.outs.assign(x)
        self._initialized.assign(True)
    else:
        self.outs.assign(x)
    
    # 基于张量运算计算更新值，保证梯度能追踪到self.w
    updated_output = self.outs + self.w
    self.outs.assign(updated_output)
    # 返回计算后的张量而非状态变量本身
    return updated_output

• 所有变量必须在__init__或build阶段创建，禁止在@tf.function内动态生成变量。
• 不要返回trainable=False的变量作为输出，必须返回参与运算的张量，确保梯度能正常传导到训练变量。
• 用tf.Variable存储初始化状态，避免Python布尔变量在图模式下被固化为常数。

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