An example of custom loss using model internals

- Tensorflow version : 2.4.0 rc1

- Colab notebook: colab.research.google.com/drive/1Hwi6auz2meKvD0ogdDSywb2E4_J1F9S_?usp=sharing

 

 

---

 

import tensorflow as tf
from tensorflow import keras

 

•  Create a custom layer calculating and adding custom loss:

class ReconLoss(keras.layers.Layer):
  def __init__(self, **kwargs):
    super().__init__(**kwargs)

  def call(self, inputs):
    x, reconstruction = inputs
    recon_loss = tf.reduce_mean(tf.square(reconstruction - x))

    self.add_loss(0.05 * recon_loss)

    return

 

•  Use the custom layer in our custom model:

class ReconstructingRegressor(keras.models.Model):
    def __init__(self, n_inputs, output_dim, **kwargs):
        super().__init__(**kwargs)
        self.hidden = [keras.layers.Dense(30, activation="selu",
                                          kernel_initializer="lecun_normal")
                       for _ in range(5)]
        self.out = keras.layers.Dense(output_dim)
        self.ReconLoss = ReconLoss()
        self.reconstruct = keras.layers.Dense(n_inputs)
       
    def call(self, inputs, training=None):
        Z = inputs
        x = inputs
        for layer in self.hidden:
            Z = layer(Z)
        reconstruction = self.reconstruct(Z)
        self.ReconLoss([x, reconstruction])
        
        return self.out(Z)

 

•  Training dummy data for test with the custom model:

dummy_x = tf.random.normal((200, 8))
dummy_y = tf.random.normal((200, 1))

model = ReconstructingRegressor(dummy_x.shape[-1], 1, dynamic=True)
model.compile(loss="mse", optimizer="nadam")
history = model.fit(dummy_x, dummy_y, epochs=2)