5

I defined a custom loss function in Keras (tensorflow backend) that is comprised of reconstruction MSE and the kullback leibler divergence between the learned probability distribution and a standard normal distribution. (It is for a variational autoencoder.)

I want to be able to slowly increase how much the cost is affected by the KL divergence term during training, with a weight called "reg", starting at reg=0.0 and increasing until it gets to 1.0. I would like the rate of increase to be tuned as a hyperparameter.(As of now, I just have the "reg" parameter set constant at 0.5.)

Is there functionality in Keras to do this?

def vae_loss(y_true,y_pred):
    reg = 0.5
    # Average cosine distance for all words in a sequence 
    reconstruction_loss = tf.reduce_mean(mean_squared_error(y_true, y_pred),1)    
    # Second part of the loss ensures the z probability distribution doesn't stray too far from normal
    KL_divergence_loss = tf.reduce_mean(tf.log(z_sigma) + tf.div((1 + tf.square(z_mu)),2*tf.square(z_sigma)) - 0.5,1)
    loss = reconstruction_loss + tf.multiply(reg,KL_divergence_loss)
    return loss
Kate A Baumli
  • 123
  • 1
  • 6
  • There's a similar implementation for custom loss in keras-contrib for the 1-Lipschitz constraint: https://github.com/keras-team/keras-contrib/blob/master/examples/improved_wgan.py#L58 – Clark Kent Aug 02 '18 at 23:19
  • This is a good example of using a variable in a loss and updating it via a callback https://stackoverflow.com/a/50975180/230462 – jonperl Jan 10 '19 at 23:58

0 Answers0