如何在自定义训练循环中实现验证损失？

Question

我一直在尝试尽早停下来研究 LSTM VAE。 在训练期间，训练损失按应有的方式计算，但验证损失为 0。 我尝试编写一个自定义 val_step 函数（类似于 train_step 但没有跟踪器）来计算损失，但我认为我无法在 vae.fit() 调用中建立该函数和 validation_data 参数之间的连接。 自定义模型类如下图：

class VAE(Model):
    def __init__(self, encoder, decoder, **kwargs):
        super(VAE, self).__init__(**kwargs)
        self.encoder = encoder
        self.decoder = decoder
        self.total_loss_tracker = tf.metrics.Mean(name="total_loss")
        self.reconstruction_loss_tracker = tf.metrics.Mean(name="reconstruction_loss")
        self.kl_loss_tracker = tf.metrics.Mean(name="kl_loss")

    def call(self, x):
        _, _, z = self.encoder(x)
        return self.decoder(z)

    @property
    def metrics(self):
        return [
            self.total_loss_tracker,
            self.reconstruction_loss_tracker,
            self.kl_loss_tracker,
        ]

    def train_step(self, data):
        with tf.GradientTape() as tape:
            z_mean, z_log_var, z = self.encoder(data)
            reconstruction = self.decoder(z)
            reconstruction_loss = tf.reduce_mean(tf.reduce_sum(losses.mse(data, reconstruction), axis=1))
            kl_loss = -0.5 * (1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var))
            kl_loss = tf.reduce_mean(tf.reduce_sum(kl_loss, axis=1))
            total_loss = reconstruction_loss + kl_loss
        grads = tape.gradient(total_loss, self.trainable_weights)
        self.optimizer.apply_gradients(zip(grads, self.trainable_weights))
        self.total_loss_tracker.update_state(total_loss)
        self.reconstruction_loss_tracker.update_state(reconstruction_loss)
        self.kl_loss_tracker.update_state(kl_loss)
        return {
            "loss": self.total_loss_tracker.result(),
            "reconstruction_loss": self.reconstruction_loss_tracker.result(),
            "kl_loss": self.kl_loss_tracker.result(),
        }

    def val_step(self, validation_data):
        _, _, z = self.encoder(validation_data)
        val_reconstruction = self.decoder(z)
        val_reconstruction_loss = tf.reduce_mean(tf.reduce_sum(losses.mse(validation_data, val_reconstruction), axis=1))
        val_kl_loss = -0.5 * (1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var))
        val_kl_loss = tf.reduce_mean(tf.reduce_sum(val_kl_loss, axis=1))
        val_total_loss = val_reconstruction_loss + val_kl_loss
        return {"total_loss": self.val_total_loss}


es = callbacks.EarlyStopping(monitor='val_total_loss',
                             mode='min',
                             verbose=1,
                             patience=5,
                             restore_best_weights=True,
                             )

vae = VAE(encoder, decoder)
vae.compile(optimizer=tf.optimizers.Adam())

vae.fit(tf_train,
        epochs=100,
        callbacks=[es],
        validation_data=tf_val,
        shuffle=True
        )

这是控制台在每个 epoch 后打印的内容（验证指标显示为 0）：

38/38 [==============================] - 37s 731ms/step - loss: 3676.8105 - reconstruction_loss: 2402.6206 - kl_loss: 149.5690 - val_total_loss: 0.0000e+00 - val_reconstruction_loss: 0.0000e+00 - val_kl_loss: 0.0000e+00

如果有人能告诉我我做错了什么，那就太好了。 提前谢谢！

更新 1： 从 val_step 定义的返回值中删除了“val_”。 有趣的是，返回调用之前的行中的 val_total_loss 是灰色的，因为它没有被使用。所以看起来这两条线之间存在脱节。

Answer 1

我认为您的代码可能会从Keras VAE example code 修改。我还努力在示例代码中添加val_loss，这是适合我的解决方案。

当validation_data 是长度小于2 的元组时，Keras 似乎会引发错误，因此我将validation_data 修改如下，

vae.fit(
    tf_train,
    epochs=100,
    callbacks=[es],
    validation_data=(valid_data,valid_data),  # <-- input X twice
    shuffle=True
)

如上修改，我们需要在收到test_step的参数后将X和y分开。还注意到返回的值是val_total_loss 而不是self.val_total_loss

def test_step(self, input_data):
    validation_data, _ = input_data # <-- Seperate X and y
    z_mean, z_log_var, z = self.encoder(validation_data)
    val_reconstruction = self.decoder(z)
    val_reconstruction_loss = tf.reduce_mean(tf.reduce_sum(losses.mse(validation_data, val_reconstruction), axis=1))
    val_kl_loss = -0.5 * (1 + z_log_var - tf.square(z_mean) - tf.exp(z_log_var))
    val_kl_loss = tf.reduce_mean(tf.reduce_sum(val_kl_loss, axis=1))
    val_total_loss = val_reconstruction_loss + val_kl_loss
    return {"total_loss": val_total_loss} # <-- modify the return value here

训练时的日志会是这样的

Epoch 00018: val_loss improved from 2304.90210 to 2304.70728, saving model to ./best_model.h5
Epoch 19/10000
31/31 [==============================] - 0s 11ms/step - loss: 2325.7858 - reconstruction_loss: 2318.3337 - kl_loss: 4.9127 - val_total_loss: 2303.8118

希望这会有所帮助:)

Answer 2

tensorflow keras fit function automatically appends "val_" 到验证损失。

尝试只返回“total_loss”，例如，

return {"total_loss": self.val_total_loss}

编辑：

另外，您正在设置 val_total_loss 但返回 self.val_total_loss

Answer 3

tensorflow 中没有val_step() API，我认为你应该使用 test_step() 代替。

https://keras.io/guides/customizing_what_happens_in_fit/