title: Tracking Multiple losses with Keras
body: <p>For networks like VAEs with competing losses, it's useful to keep track of each loss independently.  That is, it's useful to see a total loss, as well as the data term and KL-code terms.</p>
  
  <p>Is this something that is possible in Keras?  It's possible to recover the losses with vae.losses, but they are tensorflow layers, and thus can't be used in keras (eg can't create a second model that computes vae losses as output).</p>
  
  <p>It seems like a way to do this would be to add them to the metrics list on compile, but they don't fit the model of metrics.</p>
  
  <p>Here's some sample code, sorry for the length, it's mildly adapted from the example code from Keras.  The major difference is that I've explicitly moved computation of the KL div to a sampling layer, which feels more natural than the original sample code.</p>
  
  <pre><code>'''This script demonstrates how to build a variational autoencoder with Keras.
  
  Reference: "Auto-Encoding Variational Bayes" https://arxiv.org/abs/1312.6114
  '''    
  from keras.layers import Input, Dense, Lambda, Layer
  from keras.models import Model
  from keras import backend as K
  from keras import metrics
  
  batch_size = 100
  original_dim = 784
  latent_dim = 2
  intermediate_dim = 256
  epochs = 50
  epsilon_std = 1.0
  
  
  x = Input(batch_shape=(batch_size, original_dim))
  h = Dense(intermediate_dim, activation='relu')(x)
  z_mean = Dense(latent_dim)(h)
  z_log_var = Dense(latent_dim)(h)
  
  class CustomSamplingLayer(Layer):
      def __init__(self, **kwargs):
          super(CustomSamplingLayer, self).__init__(**kwargs)
  
      def kl_div_loss(self, z_mean, z_log_var):
          kl_loss = - 0.5 * K.sum(1 + z_log_var - K.square(z_mean) - K.exp(z_log_var), axis=-1)
          return K.mean(kl_loss)
  
      def call(self, inputs):
          z_mean = inputs[0]
          z_log_var = inputs[1]
          loss = self.kl_div_loss(z_mean, z_log_var)
          self.add_loss(loss, inputs=inputs)
          epsilon = K.random_normal(shape=(batch_size, latent_dim), mean=0.,
                                    stddev=epsilon_std)
          return z_mean + K.exp(z_log_var / 2) * epsilon
  
  # note that "output_shape" isn't necessary with the TensorFlow backend
  z = CustomSamplingLayer()([z_mean, z_log_var])
  
  # we instantiate these layers separately so as to reuse them later
  decoder_h = Dense(intermediate_dim, activation='relu')
  decoder_mean = Dense(original_dim, activation='sigmoid')
  h_decoded = decoder_h(z)
  x_decoded_mean = decoder_mean(h_decoded)
  
  # Custom loss layer
  class CustomVariationalLayer(Layer):
      def __init__(self, **kwargs):
          self.is_placeholder = True
          super(CustomVariationalLayer, self).__init__(**kwargs)
  
      def vae_loss(self, x, x_decoded_mean):
          xent_loss = original_dim * metrics.binary_crossentropy(x, x_decoded_mean)
          return K.mean(xent_loss)
  
      def call(self, inputs):
          x = inputs[0]
          x_decoded_mean = inputs[1]
          loss = self.vae_loss(x, x_decoded_mean)
          self.add_loss(0.0 * loss, inputs=inputs)
          return x_decoded_mean
  y = CustomVariationalLayer()([x, x_decoded_mean])
  vae = Model(x, y)
  vae.compile(optimizer='rmsprop', loss=None)
  </code></pre>