tfa.MultiOptimizer判别式分层训练与ReduceLROnPlateau回调兼容问题求解

问题根因 #

• 原生ReduceLROnPlateau默认读取优化器的全局lr属性，而tfa.optimizers.MultiOptimizer是多优化器封装，本身没有统一lr，因此抛出AttributeError
• MultiOptimizer未开放下标访问接口，其内部子优化器存储在optimizer_specs属性中，每个元素的optimizer字段对应你初始化时传入的子优化器实例，直接用下标访问会触发TypeError

方案1：修改自定义回调直接适配MultiOptimizer #

仅需修改你自定义回调中访问子优化器的代码逻辑，将下标访问替换为从optimizer_specs取子优化器即可，改造成本最低，原有训练逻辑几乎不需要变动，修改后的核心代码如下：

# 仅需修改on_epoch_end方法中子优化器访问路径即可，其余逻辑无需改动
def on_epoch_end(self, epoch, logs=None):
    logs = logs or {}
    # 取第二个优化器的学习率记录到logs中，可根据需求调整
    logs['lr'] = backend.get_value(self.model.optimizer.optimizer_specs[1].optimizer.lr)
    current = logs.get(self.monitor)
    if current is None:
      logging.warning('Learning rate reduction is conditioned on metric `%s` ' 
                      'which is not available. Available metrics are: %s',
                      self.monitor, ','.join(list(logs.keys())))

    else:
      if self.in_cooldown():
        self.cooldown_counter -= 1
        self.wait = 0

      if self.monitor_op(current, self.best):
        self.best = current
        self.wait = 0
      elif not self.in_cooldown():
        self.wait += 1
        if self.wait >= self.patience:
          # 从optimizer_specs中取出子优化器
          opt0 = self.model.optimizer.optimizer_specs[0].optimizer
          opt1 = self.model.optimizer.optimizer_specs[1].optimizer
          old_lr1 = backend.get_value(opt1.lr)
          old_lr0 = backend.get_value(opt0.lr)
          if old_lr1 > np.float32(self.min_lr):
            new_lr1 = old_lr1 * self.factor
            new_lr1 = max(new_lr1, self.min_lr)
            backend.set_value(opt1.lr, new_lr1)
            new_lr0 = old_lr0 * self.factor
            new_lr0 = max(new_lr0, self.min_lr)
            backend.set_value(opt0.lr, new_lr0)
            if self.verbose > 0:
              io_utils.print_msg(
                  f'\nEpoch {epoch +1}: ' 
                  f'ReduceLROnPlateau reducing learning rate to {new_lr0} and {new_lr1}.')
            self.cooldown_counter = self.cooldown
            self.wait = 0

方案2：官方推荐的LearningRateSchedule实现（适配性更强，推荐TPU场景使用） #

按照官方提示，自定义可动态更新的学习率调度类，为每个子优化器绑定独立的调度实例，完全符合TensorFlow设计规范，适配TPU训练场景，稳定性更高。

第一步：实现可更新的学习率调度器 #

import tensorflow as tf

class DynamicLRSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):
    def __init__(self, initial_lr, min_lr):
        super().__init__()
        self.current_lr = tf.Variable(initial_lr, trainable=False)
        self.min_lr = min_lr
    
    def __call__(self, step):
        return self.current_lr
    
    def update_lr(self, factor):
        new_lr = self.current_lr * factor
        new_lr = tf.maximum(new_lr, self.min_lr)
        self.current_lr.assign(new_lr)

第二步：初始化调度器与MultiOptimizer #

with tpu_strategy.scope():
  roberta_model = create_model(512)
  # 为不同层初始化独立的学习率调度器
  lr_sched0 = DynamicLRSchedule(initial_lr=1e-5, min_lr=1e-7)
  lr_sched1 = DynamicLRSchedule(initial_lr=1e-4, min_lr=1e-7)
  optimizers = [
        AdamWeightDecay(learning_rate=lr_sched0, weight_decay_rate=1e-5),
        AdamWeightDecay(learning_rate=lr_sched1, weight_decay_rate=1e-4)
    ]
  optimizers_and_layers = [
        (optimizers[0], roberta_model.layers[:3]),
        (optimizers[1], roberta_model.layers[3:])
    ]
  opt = tfa.optimizers.MultiOptimizer(optimizers_and_layers)
  roberta_model.compile(optimizer=opt, 
  loss=tf.keras.losses.CategoricalCrossentropy(label_smoothing=0.1), metrics=["accuracy"])

第三步：自定义回调仅触发调度器更新 #

class My_ReduceLROnPlateau(tf.keras.callbacks.Callback):
  def __init__(self, lr_schedulers, monitor='val_loss', factor=0.1, patience=10, verbose=0, mode='auto', min_delta=1e-4, cooldown=0, **kwargs):
    super().__init__()
    self.lr_schedulers = lr_schedulers
    self.monitor = monitor
    self.factor = factor
    self.patience = patience
    self.verbose = verbose
    self.min_delta = min_delta
    self.cooldown = cooldown
    # 其余初始化逻辑和_reset、in_cooldown等方法和原有实现完全一致
    self._reset()
  
  # 省略_reset、in_cooldown等复用逻辑
  
  def on_epoch_end(self, epoch, logs=None):
    logs = logs or {}
    logs['lr'] = self.lr_schedulers[1].current_lr.numpy()
    current = logs.get(self.monitor)
    if current is None:
      logging.warning('Learning rate reduction is conditioned on metric `%s` ' 
                      'which is not available. Available metrics are: %s',
                      self.monitor, ','.join(list(logs.keys())))
      return
    if self.in_cooldown():
      self.cooldown_counter -= 1
      self.wait = 0
    if self.monitor_op(current, self.best):
      self.best = current
      self.wait = 0
    elif not self.in_cooldown():
      self.wait +=1
      if self.wait >= self.patience:
        # 批量更新所有调度器的学习率
        for sched in self.lr_schedulers:
          sched.update_lr(self.factor)
        if self.verbose>0:
          lr_list = [s.current_lr.numpy() for s in self.lr_schedulers]
          print(f'\nEpoch {epoch+1}: ReduceLROnPlateau reducing learning rate to {lr_list}')
        self.cooldown_counter = self.cooldown
        self.wait = 0

第四步：实例化回调并训练 #

reduce_lr = My_ReduceLROnPlateau(lr_schedulers=[lr_sched0, lr_sched1], patience=5, min_delta=1e-4, verbose=1)
history=roberta_model.fit(train, epochs=50, validation_data=val, callbacks=[reduce_lr])

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