Есть ли способ сделать динамически взвешенное сложение функций потерь в тензорном потоке?

Question

В моем коде у меня есть две пользовательские функции потери. Одним из них является отрицательное логарифмическое правдоподобие (NLL), другим - дивергенция KL (KLD).

Каждый из них рассчитывается отдельно и добавляется к tf.keras.Model с использованием model.add_loss(NLL) и model.add_loss(KLD). * 1006. *

Для первых x эпох я хочу включить только NLL, но после этого я хочу, чтобы KLD постепенно увеличивал свое влияние (отжиг). Моя желаемая структура:

weight_KL {epoch 00 -> 09} == 0.0
weight_KL {epoch 10 -> 19} += 0.1
weight_KL {epoch 20 -> ..} == 1.0

Способ, который я вижу, состоит в том, что есть две функции потерь, добавленные вместе с весом, который изменяется за каждую эпоху:

total_loss = NLL + weight_KL * KLD

Есть ли способ выполнить sh это в TensorFlow 2?

Answer 1

Вы можете использовать callbacks аргумент в model.fit(). Добавьте вашу компенсацию потерь во встроенную функцию on_epoch_begin или on_epoch_end согласно требованию.

Определите свою функцию обратного вызова примерно так:

# Define the Required Callback Function
class LossCallback(tf.keras.callbacks.Callback):
  def on_epoch_begin(self, epoch, logs=None):
    if (epoch < 3):
      model.add_loss(NLL + KLD)
    else:
      model.add_loss(NLL + KLD +0.1)

И ваш model.fit() должен выглядеть следующим образом ниже -

model.fit(X, Y, epochs=4, callbacks=[LossCallback()])

У вас есть еще несколько опций в callbacks, таких как on_batch_begin, on_batch_end et c, вы можете найти их здесь .

Надеюсь, что это ответ на ваш вопрос.