自定义损失函数Gradient Tape报错修复：自编码器梯度正则化问题

问题背景 #

我正在研究自编码器，希望对比其与PCA的精度。导师要求添加一个包含解码器输出相对于 latent space 变量导数的自定义损失函数，但代码运行时报错：

ValueError: Passed in object <KerasTensor shape=(None, 27), dtype=float32, sparse=False, name=keras_tensor_267> of type 'KerasTensor', not tf.Tensor or tf.Variable or ExtensionType

附上的代码如下：

from keras.layers import Input, Dense, BatchNormalization, Activation
from keras.models import Model
from keras.regularizers import L2
import tensorflow as tf

# === 定义自编码器 ===
input_dim = scaled_dataset.shape[1]  # 数据集列数
input_layer = Input(shape=(input_dim,))  # 输入层

# 编码器
encoder1 = Dense(27, kernel_regularizer=L2(0.01))(input_layer)
encoder2 = BatchNormalization()(encoder1)
encoder3 = Activation('relu')(encoder2) # 潜空间

# 编码器模型
encoder_model = Model(input_layer, encoder3)

# 解码器
decoder1 = Dense(input_dim, activation='relu')(encoder3)

# 解码器模型
decoder_model = Model(encoder3, decoder1)

# 定义自编码器
autoencoder = Model(input_layer, decoder1)

# === 提取编码器和解码器 ===
encoder_model = Model(input_layer, encoder3, name="Encoder")  # 编码器模型


# === 带梯度正则的损失函数 ===
def gradient_norm_regularization(y_true, y_pred):
    # 均方误差（MSE）
    mse_loss = tf.reduce_mean(tf.square(y_true - y_pred))

    with tf.GradientTape() as grad:
        a = tf.convert_to_tensor(train.values)
        grad.watch(encoder_model.output)  # 监控潜空间
        encoder_output = encoder_model(train)  # 训练数据过编码器得到潜空间变量
        salida_reconstruida = autoencoder(train)  # 获取重构输出

    # 计算解码器输出相对于潜空间变量的导数
    gradients = grad.gradient(salida_reconstruida, encoder_output)

    # 梯度的L2范数
    norm = tf.reduce_sum(tf.square(gradients))

    # 总损失（MSE + 梯度正则）
    total_loss = mse_loss + 0.01 * norm  # 0.01是正则系数

    return total_loss

# === 编译自编码器 ===
autoencoder.compile(optimizer='adam', loss=gradient_norm_regularization)

核心报错原因 #

1. encoder_model.output是Keras的符号张量（KerasTensor），不是tf.GradientTape可监控的实际张量/变量，watch方法仅支持Tensor、Variable或扩展类型。
2. 损失函数硬编码使用全局train数据集，违反Keras设计逻辑——损失函数应基于当前批次的输入/输出，而非固定数据集。
3. 梯度计算未正确追踪潜空间变量的生成过程，导致无法求解解码器输出对其的导数。

修复步骤 #

1. 重构自编码器，同时输出重构结果和潜空间变量 #

修改模型定义，让自编码器同时返回重构输出和潜空间张量，方便损失函数获取当前批次的潜空间数据：

# 修改自编码器定义，输出[重构结果, 潜空间变量]
autoencoder = Model(input_layer, [decoder1, encoder3])

2. 重写自定义损失函数，基于当前批次计算梯度 #

利用模型的多输出特性，在损失函数中接收潜空间变量，并用GradientTape追踪其变化：

def gradient_norm_regularization(y_true, y_pred):
    # y_pred是列表：[重构输出, 潜空间变量]
    reconstructed_output, latent = y_pred
    # 计算MSE损失
    mse_loss = tf.reduce_mean(tf.square(y_true - reconstructed_output))

    with tf.GradientTape() as tape:
        # 监控当前批次的潜空间张量
        tape.watch(latent)
        # 通过解码器计算重构输出（复用已定义的decoder_model）
        decoded = decoder_model(latent)
    
    # 求解解码器输出对潜空间变量的梯度
    gradients = tape.gradient(decoded, latent)
    # 计算梯度的L2范数
    grad_norm = tf.reduce_sum(tf.square(gradients))
    
    # 总损失 = MSE + 正则项
    return mse_loss + 0.01 * grad_norm

3. 调整模型编译与训练逻辑 #

因为模型有两个输出，需要为第二个输出（潜空间变量）指定一个无意义的“占位损失”（不会影响训练）：

# 编译模型：第一个输出用自定义损失，第二个输出用dummy损失（返回0，不参与计算）
autoencoder.compile(
    optimizer='adam',
    loss=[gradient_norm_regularization, lambda y_true, y_pred: 0.0]
)

# 训练时，标签对应两个输出：第一个是原始输入（重构目标），第二个可以是任意同形状张量
autoencoder.fit(
    x=scaled_dataset,
    y=[scaled_dataset, tf.zeros((len(scaled_dataset), 27))],
    epochs=50,
    batch_size=32
)

4. 优化点：避免冗余计算 #

原代码中y_pred已经是重构结果，无需再调用autoencoder(train)重复计算，直接复用y_pred中的重构输出即可，减少计算开销。

关键原理 #

• tf.GradientTape只能追踪实际参与前向传播的张量，无法处理Keras的符号张量，因此必须使用当前批次的潜空间输出（实际Tensor）作为监控对象。
• Keras多输出模型是传递中间张量到损失函数的标准方式，避免在损失函数中硬编码数据集或模型调用，保证训练的批次一致性。

内容的提问来源于stack exchange，提问作者Jonathan Jesus Calderon Rivera