gradient tape for custom loss function

Question

I'm currently working with an autoencoder in hopes to test its accuracy vs pca. My tutor asked me to add a custom loss function that involves the derivatives of the decoder output with respect to the latent space variables.

So far I have this code but I can't get it to work. I receive errors like: 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.

So I hope I can get insight on what to do to fix this.

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

# === Definir el Autoencoder ===
input_dim = scaled_dataset.shape[1]  # Número de columnas del dataset
input_layer = Input(shape=(input_dim,))  # Capa de entrada

# Encoder
encoder1 = Dense(27, kernel_regularizer=L2(0.01))(input_layer)
encoder2 = BatchNormalization()(encoder1)
encoder3 = Activation('relu')(encoder2) #Espacio latente

#Modelo encoder
encoder_model = Model(input_layer, encoder3)

# Decoder
decoder1 = Dense(input_dim, activation='relu')(encoder3)

#Modelo decoder
decoder_model = Model(encoder3, decoder1)

# Definir el autoencoder
autoencoder = Model(input_layer, decoder1)

# === Extraer Encoder y Decoder ===
encoder_model = Model(input_layer, encoder3, name="Encoder")  # Modelo del encoder


# === Función de Pérdida con Regularización del Gradiente ===
def gradient_norm_regularization(y_true, y_pred):
    # Error cuadrático medio (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)  # Monitorear el espacio latente
        encoder_output = encoder_model(train)  # Pasar datos reales por el encoder para obtener las variables en el espacio latente
        salida_reconstruida = autoencoder(train)  # Se obtiene la salida reconstruida

    # Cálculo de la derivada de la salida del decodificador con respecto al espacio latente
    gradients = grad.gradient(salida_reconstruida, encoder_output)

    # Norma L2 de los gradientes
    norm = tf.reduce_sum(tf.square(gradients))

    # Pérdida total (MSE + regularización del gradiente)
    total_loss = mse_loss + 0.01 * norm  # 0.01 es el coeficiente de regularización

    return total_loss

# === Compilar el Autoencoder con la nueva función de pérdida ===
autoencoder.compile(optimizer='adam', loss=gradient_norm_regularization)

Answer 1

I think your issue is related to Tensorflow's GradientTape usage inside the custom loss function. Issues in your code:

1.Tensoflow's GradientTape is meant to be used within a training loop, in your code you are using it inside the loss function, which keras does not support during compilation.

2. keras loss function only accept y_true and y_pred but you function refernces train, which is not provided during model training. Here is a fixed code 1st part of code

2nd part of code

What changed &why it works:

1. Moved GradientTape to separate function(compute_gradient_norm).

2.Ensured y_true and y_pred are only used inside the loss function.

3.Used tf.function for better performance.

4. Fixed incorrect usage of train iinside loss function.

5. Added a working training example for testing