將 Tensorflow 2.0 中的內部損失函數記錄到 tensorboard

Question

我一直在為我正在使用的自定義應用程序實現我的第一個 Variational Autoencoder。一切運行順利，但我無法在每批前向傳遞時將 KL 術語記錄到 Tensorboard 。 我希望這樣做是為了深入了解模型並研究 KL 退火技術的效果。

我已經根據 TF 文檔網頁上的示例定義了我的模型，這意味着Model.add_loss()函數將 KL 項添加到當前批次損失中：

class VariationalAutoEncoder(tf.keras.Model):
def __init__(self,
             original_dim=18,
             intermediate_dim=64,
             latent_dim=2,
             name='Autoencoder',
             **kwargs):
    super(VariationalAutoEncoder, self).__init__(name=name, **kwargs)
    self.original_dim = original_dim
    self.encoder = Encoder(latent_dim=latent_dim,
                           intermediate_dim=intermediate_dim)
    self.decoder = Decoder(original_dim=original_dim,
                           intermediate_dim=intermediate_dim)
    self.kl_loss_weight = 0

def call(self, inputs):
    # Forward pass of the Encoder
    z_mean, z_log_var, z = self.encoder(inputs)
    # Forward pass of the Decoder taken the re-parameterized z latent variable
    reconstructed = self.decoder(z)

    # Add KL divergence regularization loss for this forward pass
    # THIS IS THE VALUE I WANT TO LOG TO TENSORBOARD
    kl_loss = - 0.5 * tf.reduce_mean(z_log_var - tf.square(z_mean) - tf.exp(z_log_var) + 1)

    # HERE !!!! I ADD THE KL TERM weighted by an internal variable.
    self.add_loss(self.kl_loss_weight * kl_loss)
    return reconstructed

我使用內置函數定義、編譯和訓練我的模型（我沒有定義自定義訓練循環）：

vae = VariationalAutoEncoder(original_dim=18, intermediate_dim=intermediate_dim, latent_dim=latent_dim)

optimizer = tf.keras.optimizers.Adam(learning_rate=lr)

vae.compile(optimizer,
            loss=tf.keras.losses.MeanSquaredError(),
            metrics=[tf.keras.metrics.MeanAbsoluteError(name="reconstruction_MAE")])

通過此設置，我能夠獲得組合損失函數（重建 + KL）的 Tensorboard 日志，並且也只能獲得重建（通過使用 MAE 指標）。 盡管如此，我似乎找不到同時記錄 KL 項的方法，因為它的值是在 graph 內計算的。

我嘗試使用多輸出配置，但它需要分配常量loss_weights ，因為我想使用 KL 退火，所以我需要避免這種情況。

我想避免編寫自定義訓練循環，所以我能夠做到這一點，但到目前為止我已經嘗試了幾個想法，現在我相信這可能是唯一的方法。 我錯了嗎？ 有沒有辦法將這個內部損失值記錄到 Tensorboard？

Answer 1

考慮使用tf.summary.scalar ：

class VariationalAutoEncoder(tf.keras.Model):
    <...>
    def call(self, inputs):
        <...>
        kl_loss = - 0.5 * tf.reduce_mean(z_log_var - tf.square(z_mean) - tf.exp(z_log_var) + 1)
        tf.summary.scalar('KL loss', kl_loss)
        self.add_loss(self.kl_loss_weight * kl_loss)
        return reconstructed

logdir = "logs/scalars/" + datetime.now().strftime("%Y%m%d-%H%M%S")
file_writer = tf.summary.create_file_writer(logdir + "/metrics")
file_writer.set_as_default()
tensorboard_callback = keras.callbacks.TensorBoard(log_dir=logdir)
vae = VariationalAutoEncoder(original_dim=18, intermediate_dim=intermediate_dim, latent_dim=latent_dim)

optimizer = tf.keras.optimizers.Adam(learning_rate=lr)

vae.compile(optimizer,
            loss=tf.keras.losses.MeanSquaredError(),
            metrics=[tf.keras.metrics.MeanAbsoluteError(name="reconstruction_MAE"),
            callbacks=[tensorboard_callback]