How to take the trainable parameters into a loss function in Tensoflow.Keras

Question

I'm trying to implement a loss function where variables in Convolutional layers are required for calculation. There's one method given by the official documents that involve variables in the loss function:

If this is not the case for your loss (if, for example, your loss references a Variable of one of the model's layers), you can wrap your loss in a zero-argument lambda. These losses are not tracked as part of the model's topology since they can't be serialized.

inputs = tf.keras.Input(shape=(10,))
x = tf.keras.layers.Dense(10)(inputs)
outputs = tf.keras.layers.Dense(1)(x)
model = tf.keras.Model(inputs, outputs)
# Weight regularization.
model.add_loss(lambda: tf.reduce_mean(x.kernel))

However, this is just adding a simple regularized to the model. Is there a way to implement more complicated regularizer which involves calculation between variables in different layers? And what if a trainable variable is added to the regularizer as well?

Answer 1

You can add arbitarily complex loss functions using the add_loss API. Here is an example to add a loss that uses weights of two different layers.

import tensorflow as tf

print('TensorFlow:', tf.__version__)

inp = tf.keras.Input(shape=[10])
x = tf.keras.layers.Dense(16)(inp)
x = tf.keras.layers.Dense(32)(x)
x = tf.keras.layers.Dense(4)(x)
out = tf.keras.layers.Dense(1)(x)

model = tf.keras.Model(inputs=[inp], outputs=[out])
model.summary()


def custom_loss(weight_a, weight_b):
    def _custom_loss():
        # This can include any arbitrary logic
        loss = tf.norm(weight_a) + tf.norm(weight_b)
        return loss
    return _custom_loss

weight_a = model.layers[2].kernel
weight_b = model.layers[3].kernel

model.add_loss(custom_loss(weight_a, weight_b))


print('\nlosses:', model.losses)

Output:

TensorFlow: 2.3.0-dev20200611
Model: "functional_1"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
input_1 (InputLayer)         [(None, 10)]              0         
_________________________________________________________________
dense (Dense)                (None, 16)                176       
_________________________________________________________________
dense_1 (Dense)              (None, 32)                544       
_________________________________________________________________
dense_2 (Dense)              (None, 4)                 132       
_________________________________________________________________
dense_3 (Dense)              (None, 1)                 5         
=================================================================
Total params: 857
Trainable params: 857
Non-trainable params: 0
_________________________________________________________________

losses: [<tf.Tensor: shape=(), dtype=float32, numpy=7.3701963>]