Callback with add_loss in Keras

user_12767924 0
Viewed 596

I am trying to implement a VAE style network in Keras. I compute my negative log-likelihood and KL divergence and add them to my model with model.add_loss(NLL) and model.add_loss(KL), respectively.

I would like to scale my KL term as training progresses ("KL-annealing"). I attempted to do this with custom callback (detailed here), but the KL loss term does not get updated - the model.add_loss(KL) term is constant over time, despite the KL weight getting updated (see Figure).

Loss over training epochs How can I make the model.add_loss(KL) term depend on KL_weight?

Code to demonstrate the idea:

...
<NLL calculations>
...

# Add the first loss to the model, the NLL:
model.add_loss(NLL)

from keras.callbacks import Callback

klstart = 40
# number of epochs over which KL scaling is increased from 0 to 1
kl_annealtime = 20

weight = tf.keras.backend.variable(0.0) #intialiase KL weight term
class AnnealingCallback(Callback):
    def __init__(self, weight):
        self.weight = weight
    def on_epoch_end (self, epoch, logs={}):
        if epoch > klstart :
            new_weight = min(tf.keras.backend.get_value(self.weight) + (1./ kl_annealtime), 1.)
            tf.keras.backend.set_value(self.weight, new_weight)
        print ("Current KL Weight is " + str(tf.keras.backend.get_value(self.weight)))

# Now the KL divergence:
<KL weight calculations>
KL = weight*tf.reduce_mean(tfp.distributions.kl_divergence(p, q))
model.add_loss(KL)

# Now compile the model with a specified optimiser
opt = tf.keras.optimizers.Adam(lr=0.001,clipnorm=0.1)
model.compile(optimizer=opt)

# Monitor how the NLL and KL divergence differ over time
model.add_metric(KL, name='kl_loss', aggregation='mean')
model.add_metric(NLL, name='mse_loss', aggregation='mean')

ops.reset_default_graph()
history=model.fit(Y_portioned, # Input or "Y_true"
                    verbose=1,
                    callbacks=[earlystopping_callback,callback_reduce_lr,AnnealingCallback(weight)],
                    epochs=650,
                    batch_size=8
                    ) # <- Increase batch size for speed up

Versions: TensorFlow 2.1.0, Keras 2.2.4

Many thanks in advance

keras tensorflow deep-learning machine-learning python
1 Answers

So I was exactly looking for this implementation in Keras when I came across your post. After a while I seem to have found only two mistakes:

  1. You declared a tf variable out of the class context so it was not recognized.
  2. for the first reason you didn't properly instantiate the weight.

the corrected version would look like this:

#start = klstart
#time = kl_annealtime

class AnnealingCallback(keras.callbacks.Callback):
    def __init__(self, weight=tf.keras.backend.variable(0.0), start=20, time=40):
        self.weight = weight
        self.start = start
        self.time = time    
    def on_epoch_end (self, epoch, logs={}):
        if epoch > self.start :
            new_weight = min(tf.keras.backend.get_value(self.weight) + (1./self.time), 1.)
            tf.keras.backend.set_value(self.weight, new_weight)
        print("Current KL Weight is " + str(tf.keras.backend.get_value(self.weight)))

Now you can instantiate the weight:

AC = AnnealingCallback()
w = AC.weight

which pre-multpiplies your KL-divergence.

user_152061970
Related