为训练和验证指定单独的 tf.summary file_writers

Posted 2023-03-27

【中文标题】为训练和验证指定单独的 tf.summary file_writers

【英文标题】：Specifying separate tf.summary file_writers for training and validation

【发布时间】：2021-03-10 14:12:03

【问题描述】：

随着训练的进行，我想可视化一些标量，例如 Train loss、val loss（以及其他）。我将 tf.Keras 与 tf2.3 一起使用。但是，我无法为训练和验证指定单独的文件编写器，这会导致 tensorboard 中的图形全部损坏。 我的张量板如下所示： Tensorboard output 我还在每次训练运行前清理日志。所以这不是以前的运行。 问题是我只能设置一个默认作家。那么如何根据当前是通过损失函数运行训练还是评估来切换编写器呢？

伪代码：

def loss_fn():
  ..calculate loss..
  tf.summary.scalar('loss', loss)
def train():
  writer = tf.summary.create_file_writer(os.path.join(args.training_folder, 'logs')
  writer.set_as_default()
  model = create_model()
  model.compile(*arguments here*)
  model.fit(*arguments here*)

【参考方案1】：

修复很简单。这可以通过在 keras 回调的 on_train_begin() 和 on_test_begin() 方法中设置不同的编写器来解决

class TensorBoardFix(tf.keras.callbacks.TensorBoard):
    
    def __init__(self, training_folder, **kwargs):
        super().__init__()
        self.train_writer = tf.summary.create_file_writer(os.path.join(training_folder, 'logs', 'train'))
        self.val_writer = tf.summary.create_file_writer(os.path.join(training_folder, 'logs', 'val'))

    def on_train_begin(self, *args, **kwargs):
        super(TensorBoardFix, self).on_train_begin(*args, **kwargs)
        tf.summary.experimental.set_step(self._train_step)
        self.train_writer.set_as_default()

    def on_test_begin(self, *args, **kwargs):
        super(TensorBoardFix, self).on_test_begin(*args, **kwargs)
        tf.summary.experimental.set_step(self._val_step)
        self.val_writer.set_as_default()