python - 在 Keras 中使用 add_loss 回调

Question

我正在尝试在 Keras 中实现 VAE 风格的网络。我计算负对数似然和 KL 散度，并分别使用 model.add_loss(NLL) 和 model.add_loss(KL) 将它们添加到我的模型中。

我想随着训练的进行扩展我的 KL 术语(“KL 退火”)。我尝试使用自定义回调 ( detailed here ) 执行此操作，但 KL 损失项未更新 - model.add_loss(KL) 项随时间保持不变，尽管 KL 权重已更新(见图)

Loss over training epochs 如何使 model.add_loss(KL) 项依赖于 KL_weight？

演示这个想法的代码:

...
<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

版本:TensorFlow 2.1.0、Keras 2.2.4

提前致谢

Answer 1

所以当我看到你的帖子时，我正在寻找 Keras 中的这个实现。过了一会儿我似乎只发现了两个错误:

1. 您在类上下文之外声明了一个 tf 变量，因此无法识别它。
2. 第一个原因是您没有正确实例化权重。

修正后的版本应该是这样的:

#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)))

现在你可以实例化权重了:

AC = AnnealingCallback()
w = AC.weight

它预先乘以你的 KL 散度。