通过权重和偏差在 keras 中记录学习率计划

Question

我正在训练 keras 模型并为优化器使用自定义学习率调度程序（类型为 tf.keras.optimizers.schedules.LearningRateSchedule），并且我想通过权重和偏差框架记录学习率变化。我找不到如何将其传递给 WandbCallback 对象或以任何方式记录它

Answer 1

根据 Martjin 的评论进行更新！

您可以使用自定义 Keras 回调将自定义学习率记录到权重和偏差中。

W&BWandbCallback无法自动记录您的自定义学习率。通常，对于此类自定义日志记录，如果您使用自定义训练循环，则可以使用wandb.log(). 如果您使用model.fit()自定义 Keras 回调，就是这样。

例如：

这是我的tf.keras.optimizers.schedules.LearningRateSchedule基于调度程序。

class MyLRSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):

  def __init__(self, initial_learning_rate):
    self.initial_learning_rate = initial_learning_rate

  def __call__(self, step):
     return self.initial_learning_rate / (step + 1)

optimizer = tf.keras.optimizers.SGD(learning_rate=MyLRSchedule(0.001))

您可以使用 获取优化器的当前学习率optimizer.learning_rate(step)。这可以包装为自定义 Keras 回调并wandb.log()与其一起使用。

class LRLogger(tf.keras.callbacks.Callback):
    def __init__(self, optimizer):
      super(LRLogger, self).__init__()
      self.optimizer = optimizer

    def on_epoch_end(self, epoch, logs):
      lr = self.optimizer.learning_rate(self.optimizer.iterations)
      wandb.log({"lr": lr}, commit=False)

请注意，在wandb.log通话中我使用了commit=False参数。这将确保在同一时间步骤记录每个指标。更多内容请参见此处。

称呼model.fit()。

tf.keras.backend.clear_session()
model = some_model()

model.compile(optimizer, 'categorical_crossentropy', metrics=['acc'])

wandb.init(entity='wandb-user-id', project='my-project', job_type='train')

_ = model.fit(trainloader,
          epochs=EPOCHS,
          validation_data=testloader,
          callbacks=[WandbCallback(), # using WandbCallback to log default metrics.
                     LRLogger(optimizer)]) # using callback to log learning rate.

wandb.finish()

这是 W&B 媒体面板：