Keras反馈循环自编码器中TensorArray导致无梯度的问题求助

问题重现 #

以下是实现反馈循环自编码器的代码：

import tensorflow as tf
import keras
import os

class Linear(keras.layers.Layer):
    def __init__(self, units=32):
        super(Linear, self).__init__()
        self.units = units

    def build(self, input_shape):
        self.w = self.add_weight(
            shape=(input_shape[-1], self.units),
            initializer="random_normal",
            trainable=True,
        )
        self.b = self.add_weight(
            shape=(self.units,), initializer="random_normal", trainable=True
        )

    def call(self, inputs):
        return tf.matmul(inputs, self.w) + self.b


class FRAE(tf.keras.Model):
    def __init__(self):
        super(FRAE, self).__init__()
        self.linear_1 = Linear(4)
        self.linear_2 = Linear(3)
        self.latent   = Linear(1)
        self.linear_3 = Linear(3)
        self.linear_4 = Linear(2)
        self.decoded  = tf.zeros(shape=(1, 2))
        
    def call(self, inputs):
        batch_size = inputs.shape[0]
        output_list = tf.TensorArray(tf.float32, size=batch_size, clear_after_read=False)

        for i in range(batch_size):
            x = tf.concat((tf.expand_dims(inputs[i], axis=0),self.decoded),axis=1)
            x = self.linear_1(x)
            x = tf.nn.swish(x)
            x = self.linear_2(x)
            x = tf.nn.swish(x)
            x = self.latent(x)
            x = tf.nn.swish(x)
            x = tf.concat((x,self.decoded),axis=1)
            x = self.linear_3(x)
            x = tf.nn.swish(x)
            x = self.linear_4(x)
            x = tf.nn.swish(x)
            self.decoded = tf.identity(x)
            output_list.write(i,  x)
        y = output_list.stack()
        return y


os.environ['CUDA_VISIBLE_DEVICES'] = '-1'

xtrain = tf.random.uniform(shape=(10,2))
model = FRAE()
y = model(xtrain)
optimizer = keras.optimizers.Adam(lr=0.001)
model.compile(optimizer=optimizer,loss="mse")
model.fit(x=xtrain,y=xtrain, epochs=50, batch_size=1)

运行后出现梯度丢失错误：

ValueError: No gradients provided for any variable: (['frae_13/linear_65/Variable:0', 'frae_13/linear_65/Variable:0', 'frae_13/linear_66/Variable:0', 'frae_13/linear_66/Variable:0', 'frae_13/linear_67/Variable:0', 'frae_13/linear_67/Variable:0', 'frae_13/linear_68/Variable:0', 'frae_13/linear_68/Variable:0', 'frae_13/linear_69/Variable:0', 'frae_13/linear_69/Variable:0'],). Provided grads_and_vars is ((None, <tf.Variable 'frae_13/linear_65/Variable:0' shape=(4, 4) dtype=float32>), (None, <tf.Variable 'frae_13/linear_65/Variable:0' shape=(4,) dtype=float32>), (None, <tf.Variable 'frae_13/linear_66/Variable:0' shape=(4, 3) dtype=float32>), (None, <tf.Variable 'frae_13/linear_66/Variable:0' shape=(3,) dtype=float32>), (None, <tf.Variable 'frae_13/linear_67/Variable:0' shape=(3, 1) dtype=float32>), (None, <tf.Variable 'frae_13/linear_67/Variable:0' shape=(1,) dtype=float32>), (None, <tf.Variable 'frae_13/linear_68/Variable:0' shape=(3, 3) dtype=float32>), (None, <tf.Variable 'frae_13/linear_68/Variable:0' shape=(3,) dtype=float32>), (None, <tf.Variable 'frae_13/linear_69/Variable:0' shape=(3, 2) dtype=float32>), (None, <tf.Variable 'frae_13/linear_69/Variable:0' shape=(2,) dtype=float32>)).

错误原因 #

1. 类属性存储状态导致梯度断裂：self.decoded作为类属性存储循环状态，在TensorFlow的计算图模式下（Kerasfit会自动启用），这种方式无法被梯度追踪机制识别，导致后续操作的梯度无法回溯到模型变量。
2. Python循环与计算图兼容性问题：直接使用Pythonfor循环处理批量样本，在静态计算图中无法正确构建完整的梯度路径，加上TensorArray的使用方式没有配合可追踪的状态传递，进一步加剧了梯度丢失。

解决方案 #

改用tf.while_loop构建可追踪的循环逻辑，将状态变量（decoded）作为循环的中间张量传递，而非存储为类属性，确保梯度能完整回溯：

import tensorflow as tf
import keras
import os

class Linear(keras.layers.Layer):
    def __init__(self, units=32):
        super(Linear, self).__init__()
        self.units = units

    def build(self, input_shape):
        self.w = self.add_weight(
            shape=(input_shape[-1], self.units),
            initializer="random_normal",
            trainable=True,
        )
        self.b = self.add_weight(
            shape=(self.units,), initializer="random_normal", trainable=True
        )

    def call(self, inputs):
        return tf.matmul(inputs, self.w) + self.b


class FRAE(tf.keras.Model):
    def __init__(self):
        super(FRAE, self).__init__()
        self.linear_1 = Linear(4)
        self.linear_2 = Linear(3)
        self.latent   = Linear(1)
        self.linear_3 = Linear(3)
        self.linear_4 = Linear(2)
        
    def call(self, inputs):
        batch_size = tf.shape(inputs)[0]
        # 初始化状态张量，而非类属性
        decoded = tf.zeros(shape=(1, 2))
        # 初始化TensorArray存储输出
        output_array = tf.TensorArray(tf.float32, size=batch_size)

        def loop_body(i, decoded, output_array):
            # 获取当前样本
            current_input = tf.expand_dims(inputs[i], axis=0)
            # 编码器部分
            x = tf.concat([current_input, decoded], axis=1)
            x = self.linear_1(x)
            x = tf.nn.swish(x)
            x = self.linear_2(x)
            x = tf.nn.swish(x)
            x = self.latent(x)
            x = tf.nn.swish(x)
            # 解码器部分
            x = tf.concat([x, decoded], axis=1)
            x = self.linear_3(x)
            x = tf.nn.swish(x)
            x = self.linear_4(x)
            x = tf.nn.swish(x)
            # 更新状态和输出数组
            decoded = x
            output_array = output_array.write(i, x)
            return i + 1, decoded, output_array

        # 执行while循环
        _, _, output_array = tf.while_loop(
            cond=lambda i, *_: i < batch_size,
            body=loop_body,
            loop_vars=(0, decoded, output_array),
            # 确保输出数组形状可推断
            shape_invariants=(
                tf.TensorShape([]),
                tf.TensorShape((1, 2)),
                tf.TensorShape(None)
            )
        )

        # 将TensorArray转为张量返回，调整形状匹配输入
        outputs = output_array.stack()
        outputs = tf.squeeze(outputs, axis=1)
        return outputs


os.environ['CUDA_VISIBLE_DEVICES'] = '-1'

xtrain = tf.random.uniform(shape=(10,2))
model = FRAE()
optimizer = keras.optimizers.Adam(learning_rate=0.001)
model.compile(optimizer=optimizer,loss="mse")
model.fit(x=xtrain,y=xtrain, epochs=50, batch_size=1)

关键说明 #

1. 使用tf.while_loop替代Python循环：tf.while_loop是TensorFlow原生的循环操作，能在静态计算图中完整保留梯度路径，确保所有操作都被梯度追踪机制覆盖。
2. 状态变量作为循环参数传递：将decoded作为循环的输入/输出参数，而非类属性，保证状态更新过程被记录在计算图中，梯度可以正常回溯。
3. 指定形状不变量：在tf.while_loop中设置shape_invariants，帮助TensorFlow推断循环中张量的形状，避免形状不确定导致的梯度计算失败。
4. 调整输出形状：最终将TensorArray堆叠后的张量挤压维度，确保输出形状和输入一致（(batch_size, 2)），匹配MSE损失的计算要求。

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