Output multiple losses added by add_loss in Keras

Question

I've investigated the Keras example for custom loss layer demonstrated by a Variational Autoencoder (VAE). They have only one loss-layer in the example while the VAE's objective consists out of two different parts: Reconstruction and KL-Divergence. However, I'd like to plot/visualize how these two parts evolve during training and split the single custom loss into two loss-layer:

Keras Example Model:

My Model:

Unfortunately, Keras just outputs one single loss value in the for my multi-loss example as can be seen in my Jupyter Notebook example where I've implemented both approaches. Does someone know how to get the values per loss which were added by add_loss? And additionally, how does Keras calculate the single loss value, given multiple add_loss calls (Mean/Sum/...?)?

score 11 · Answer 1 · edited Nov 25 '20 at 03:20

11

I'm using the version 2.2.4-tf of Keras and the solution above didn't work for me. Here is the solution I found (to continue the example of dumkar):

reconstruction_loss = mse(K.flatten(inputs), K.flatten(outputs))
kl_loss = beta*K.mean(- 0.5 * 1/latent_dim * K.sum(1 + z_log_var - K.square(z_mean) - K.exp(z_log_var), axis=-1))

model.add_loss(reconstruction_loss)
model.add_loss(kl_loss)

model.add_metric(kl_loss, name='kl_loss', aggregation='mean')
model.add_metric(reconstruction_loss, name='mse_loss', aggregation='mean')

model.compile(optimizer='adam')

Hope it will help you.

edited Nov 25 '20 at 03:20

Aaron Voelker

729
6
12

answered Dec 16 '19 at 18:50

Valentin Vignal

6,151
2
33
73

2

Thanks for the extension! This is, in fact, the only feasible solution in tf-keras (aka tensorflow 2.0). – Tik0 Dec 17 '19 at 09:36
1

I've implemented, just as a reference, your proposed solution for the standard VAE Keras example here: https://gist.github.com/tik0/6aa42cabb9cf9e21567c3deb309107b7 – Tik0 Dec 17 '19 at 11:31
@ValentinVIGNAL Thank you for the answer. can you please tell me (curious and want to learn) how you understood/thought of passing the aggregation ='mean' argument to the add_metric. I don't see any where any example that uses other than the function, name arguments. There are **kwargs but not sure what they are referring to. – tjt Sep 15 '20 at 03:04
1

Using a newer version of TF (2.2.0), you need to move the `model.compile` call after the `model.add_metric` calls, otherwise you get a `TypeError: unhashable type: 'DictWrapper'` error (I've proposed this as an edit to this post). – Aaron Voelker Nov 24 '20 at 16:46

dumkar · Accepted Answer · 2019-03-06T14:57:45.510

This is indeed not supported, and currently discussed on different places on the web. The solution can be obtained by adding your losses again as a separate metric after the compile step (also discussed here)

This results in something like this (specifically for a VAE):

reconstruction_loss = mse(K.flatten(inputs), K.flatten(outputs))
kl_loss = beta*K.mean(- 0.5 * 1/latent_dim * K.sum(1 + z_log_var - K.square(z_mean) - K.exp(z_log_var), axis=-1))

model.add_loss(reconstruction_loss)
model.add_loss(kl_loss)
model.compile(optimizer='adam')

model.metrics_tensors.append(kl_loss)
model.metrics_names.append("kl_loss")

model.metrics_tensors.append(reconstruction_loss)
model.metrics_names.append("mse_loss")

For me this gives an output like this:

Epoch 1/1
252/252 [==============================] - 23s 92ms/step - loss: 0.4336 - kl_loss: 0.0823 - mse_loss: 0.3513 - val_loss: 0.2624 - val_kl_loss: 0.0436 - val_mse_loss: 0.2188

score 1 · Answer 3 · answered Jan 09 '19 at 11:53

It turns out that the answer is not straight forward and furthermore, Keras does not support this feature out of the box. However, I've implemented a solution where each loss-layer outputs the loss and a customized callback function records it after every epoch. The solution for my multi-headed example can be found here: https://gist.github.com/tik0/7c03ad11580ae0d69c326ac70b88f395

Output multiple losses added by add_loss in Keras

3 Answers3