TensorFlow 2.10两种学习率调度实现差异及问题修复

一、两种学习率调度方法的核心差异 #

tf.keras.optimizers.schedules.LearningRateSchedule #

• 属于优化器层级的调度逻辑，直接与优化器绑定，会在**每个训练步骤（step）**更新学习率
• 使用时需将调度规则封装为可调用对象，传入优化器的learning_rate参数，示例代码：

  lr_schedule = tf.keras.optimizers.schedules.ExponentialDecay(
      initial_learning_rate=1e-3,
      decay_steps=10000,
      decay_rate=0.9
  )
  optimizer = tf.keras.optimizers.Adam(learning_rate=lr_schedule)
  

• 支持细粒度的step级调整，适合需要根据训练步数动态调参的场景

tf.keras.callbacks.LearningRateScheduler #

• 属于回调层级的调度，通过回调函数在每个epoch结束后更新学习率
• 只需定义一个接收epoch和当前lr的函数，传入回调即可，示例代码：

  def lr_schedule_func(epoch, lr):
      if epoch < 10:
          return lr
      else:
          return lr * 0.95
  lr_callback = tf.keras.callbacks.LearningRateScheduler(lr_schedule_func)
  

• 实现简单，适合按epoch维度调整学习率的场景

二、model.fit(verbose=1)崩溃问题修复 #

崩溃原因大概率是自定义LR_History回调在读取学习率时，与verbose=1的实时日志打印操作产生资源/线程冲突，GPU环境下更容易触发这类问题。

修复方案 #

1. 规范学习率获取逻辑：避免直接访问优化器内部私有变量，改用兼容两种调度模式的安全读取方式：

   class LR_History(tf.keras.callbacks.Callback):
       def on_train_begin(self, logs=None):
           self.lrs = []
   
       def on_epoch_end(self, epoch, logs=None):
           # 区分两种调度模式，安全获取当前学习率
           lr_obj = self.model.optimizer.lr
           if isinstance(lr_obj, tf.keras.optimizers.schedules.LearningRateSchedule):
               # 针对LearningRateSchedule，传入当前训练步数计算
               current_step = self.model.optimizer.iterations.numpy()
               current_lr = lr_obj(current_step).numpy()
           else:
               # 针对固定学习率或LearningRateScheduler的情况
               current_lr = lr_obj.numpy()
           self.lrs.append(current_lr)
   

2. GPU环境额外优化：如果是显存或线程冲突导致的崩溃，可在训练前添加显存增长配置：

   gpus = tf.config.list_physical_devices('GPU')
   if gpus:
       tf.config.experimental.set_memory_growth(gpus[0], True)
   

3. 精简回调逻辑：确保自定义回调仅专注于学习率记录，不要在回调中执行其他可能干扰训练流程的操作

内容的提问来源于stack exchange，提问作者mkdelete