[FIXED] How to refer to ground truth variables in tensorflow keras?

Issue

I'm trying to add a custom loss to my model and need a reference to the target variables to do this. E.g.

model = Model(inputs=[x1, x2, x3], outputs=[y1,y2,y3])
mse = tf.keras.losses.MeanSquaredError()
model.add_loss(mse(x1, dec_x1))

Here I added a loss between one input variable and an encoded and then decoded version of the same variable. But I would also like to be able to add a loss which depends on the ground truth of the y variables (not the predicted y1, y2, y3) I.e. imagine y1_true and then to add a loss:

# Code to make Y1 which depends on x1, x2, x3
model.add_loss(mse(Y1, y1_true))

But how do I access y1_true in keras?

Solution

Unfortunately, the add_loss function does not have access to y1_true labels. This method is actually intended for regularization losses in your model:

When writing the call method of a custom layer or a subclassed model, you may want to compute scalar quantities that you want to minimize during training (e.g. regularization losses). You can use the add_loss() layer method to keep track of such loss terms.

See the docs for more information. I think you would be far better off using a custom training loop, where you would have direct access to everything you need. Here is a simplified example:

class Autoencoder(tf.keras.Model):
  def __init__(self, latent_dim):
    super(Autoencoder, self).__init__()

    self.latent_dim = latent_dim
    self.dense1 = tf.keras.layers.Dense(self.latent_dim, activation='relu')
    self.dense2 = tf.keras.layers.Dense(5, activation='relu')
    e_input1 = tf.keras.Input(shape=(5,))
    e_input2 = tf.keras.Input(shape=(5,))
    e_input3 = tf.keras.Input(shape=(5,))
    e_output1 = self.dense1(e_input1)
    e_output2 = self.dense1(e_input2)
    e_output3 = self.dense1(e_input3)

    self.encoder = tf.keras.Model([e_input1, e_input2, e_input3], [e_output1, e_output2, e_output3])

    d_input1 = tf.keras.Input(shape=(self.latent_dim,))
    d_input2 = tf.keras.Input(shape=(self.latent_dim,))
    d_input3 = tf.keras.Input(shape=(self.latent_dim,))

    d_output1 = self.dense2(d_input1)
    d_output2 = self.dense2(d_input2)
    d_output3 = self.dense2(d_input3)
    
    self.decoder = tf.keras.Model([d_input1, d_input2, d_input3], [d_output1, d_output2, d_output3])

  def encode(self, inputs):
    x1, x2, x3 = inputs
    return self.encoder([x1, x2, x3])
  
  def decode(self, inputs):
    x1, x2, x3 = inputs
    return self.decoder([x1, x2, x3])


latent_dim = 5 
autoencoder = Autoencoder(latent_dim) 
optimizer = tf.keras.optimizers.Adam()
mse = tf.keras.losses.MeanSquaredError()


your_train_dataset = tf.data.Dataset.from_tensor_slices((tf.random.normal((4, 5)), 
                                                         tf.random.normal((4, 5)), 
                                                         tf.random.normal((4, 5)), 
                                                         tf.random.normal((4, 5)),
                                                         tf.random.normal((4, 5)),
                                                         tf.random.normal((4, 5)))).batch(2)
epochs = 2
for epoch in range(epochs):
    for batch in your_train_dataset:
        x1_batch_train, x2_batch_train, x3_batch_train, y1_batch_train, y2_batch_train, y3_batch_train = batch
        with tf.GradientTape() as tape:
            enc_x1, enc_x2, enc_x1 = autoencoder.encode([x1_batch_train, x2_batch_train, x3_batch_train])
            dec_x1, dec_x2, dec_x1 = autoencoder.decode([enc_x1, enc_x2, enc_x1])

            loss1 = mse(x1_batch_train, enc_x1)
            loss2 = mse(x1_batch_train, dec_x1)
            loss3 = mse(dec_x1, y1_batch_train)
            #..... and so on.

            losses = loss1 + loss2 + loss3
            tf.print(losses)
        grads = tape.gradient(losses, autoencoder.trainable_weights)
        optimizer.apply_gradients(zip(grads, autoencoder.trainable_weights))

Answered By - AloneTogether