使用RoBERTa时添加Keras回调出现WarmUp与int相乘TypeError

问题描述 #

使用RoBERTa构建文本分类模型时，添加ReduceLROnPlateau或LearningRateScheduler回调后，触发以下类型错误：

TypeError: unsupported operand type(s) for *: 'WarmUp' and 'int'

相关代码 #

epochs = 30
steps_per_epoch = tf.data.experimental.cardinality(train_ds).numpy()
num_train_steps = steps_per_epoch * epochs
num_warmup_steps = int(0.1 * num_train_steps)
init_lr = 3e-5

callback = [tf.keras.callbacks.EarlyStopping(monitor='val_accuracy', 
                                             min_delta=0,
                                             patience=3,
                                             verbose=1,
                                             mode='auto',
                                             baseline=None,
                                             restore_best_weights=False,
                                             start_from_epoch=0),
           tf.keras.callbacks.ReduceLROnPlateau(monitor='val_loss', 
                                                factor=0.2,
                                                patience=5, 
                                                min_lr=0.001)
           ]

optimizer = optimization.create_optimizer(init_lr=init_lr,
                                          num_train_steps=num_train_steps,
                                          num_warmup_steps=num_warmup_steps,
                                          optimizer_type='adamw')

classifier_model.compile(optimizer=optimizer,
                         loss=loss,
                         metrics=metrics)

history = classifier_model.fit(x=train_ds,
                               validation_data=val_ds,
                               epochs=epochs,
                               callbacks=callback,
                               steps_per_epoch=steps_per_epoch,
                               verbose=1)

完整报错 #

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[63], line 35
     29 classifier_model.compile(optimizer=optimizer,
     30                          loss=loss,
     31                          metrics=metrics)
     33 print(f'Training model with {tfhub_handle_encoder}')
---> 35 history = classifier_model.fit(x=train_ds,
     36                                    validation_data=val_ds,
     37                                    epochs=epochs,
     38                                    callbacks=callback,
     39                                    steps_per_epoch=steps_per_epoch,
     40                                    verbose=1)

File /usr/local/lib/python3.8/site-packages/keras/utils/traceback_utils.py:70, in filter_traceback.<locals>.error_handler(*args, **kwargs)
     67     filtered_tb = _process_traceback_frames(e.__traceback__)
     68     # To get the full stack trace, call:
     69     # `tf.debugging.disable_traceback_filtering()`
---> 70     raise e.with_traceback(filtered_tb) from None
     71 finally:
     72     del filtered_tb

File /usr/local/lib/python3.8/site-packages/keras/utils/generic_utils.py:210, in Progbar.update(self, current, values, finalize)
    208 value_base = max(current - self._seen_so_far, 1)
    209 if k not in self._values:
--> 210     self._values[k] = [v * value_base, value_base]
    211 else:
    212     self._values[k][0] += v * value_base

TypeError: unsupported operand type(s) for *: 'WarmUp' and 'int'

问题原因 #

optimization.create_optimizer生成的AdamW优化器，其学习率（lr）是一个WarmUp调度对象，而非固定数值。ReduceLROnPlateau回调会尝试对优化器的lr执行乘法操作（比如乘以factor=0.2），但WarmUp对象不支持与整数/浮点数直接运算，因此触发类型错误。

解决方案 #

方法1：自定义学习率调度器，整合WarmUp与LR衰减逻辑 #

手动实现包含warmup和验证集loss监控衰减的学习率调度器，替代原有的optimization.create_optimizer和ReduceLROnPlateau回调：

import tensorflow as tf

class WarmUpWithReduceLROnPlateau(tf.keras.optimizers.schedules.LearningRateSchedule):
    def __init__(self, init_lr, num_warmup_steps, steps_per_epoch, factor=0.2, patience_epochs=5, min_lr=1e-7):
        super().__init__()
        self.init_lr = tf.Variable(init_lr, trainable=False, dtype=tf.float32)
        self.num_warmup_steps = num_warmup_steps
        self.steps_per_epoch = steps_per_epoch
        self.factor = factor
        self.patience_steps = patience_epochs * steps_per_epoch
        self.min_lr = min_lr
        
        # 用于监控验证集loss的变量
        self.best_val_loss = tf.Variable(float('inf'), trainable=False)
        self.wait_steps = tf.Variable(0, trainable=False, dtype=tf.int64)

    def __call__(self, step):
        # Warmup阶段：线性递增到初始学习率
        if step < self.num_warmup_steps:
            warmup_lr = tf.cast(step, tf.float32) / tf.cast(self.num_warmup_steps, tf.float32) * self.init_lr
            return warmup_lr
        # Warmup结束后，返回当前设置的学习率
        else:
            return self.init_lr

    def update_lr_based_on_val_loss(self, val_loss):
        # 更新最佳loss和等待步数
        if val_loss < self.best_val_loss:
            self.best_val_loss.assign(val_loss)
            self.wait_steps.assign(0)
        else:
            self.wait_steps.assign_add(1)
            # 达到耐心步数时，衰减学习率
            if self.wait_steps >= self.patience_steps:
                new_lr = self.init_lr * self.factor
                self.init_lr.assign(tf.maximum(new_lr, self.min_lr))
                self.wait_steps.assign(0)
                print(f"学习率衰减至: {self.init_lr.numpy():.6f}")

# 自定义回调，用于在每个epoch结束后更新学习率
class LRSchedulerCallback(tf.keras.callbacks.Callback):
    def __init__(self, lr_schedule):
        super().__init__()
        self.lr_schedule = lr_schedule
    
    def on_epoch_end(self, epoch, logs=None):
        val_loss = logs.get('val_loss')
        if val_loss is not None:
            self.lr_schedule.update_lr_based_on_val_loss(val_loss)

替换原有优化器和回调 #

# 初始化自定义学习率调度器
lr_schedule = WarmUpWithReduceLROnPlateau(
    init_lr=init_lr,
    num_warmup_steps=num_warmup_steps,
    steps_per_epoch=steps_per_epoch,
    factor=0.2,
    patience_epochs=5,
    min_lr=1e-7
)

# 创建AdamW优化器，传入自定义调度器
optimizer = tf.keras.optimizers.AdamW(learning_rate=lr_schedule)

# 替换回调列表
callback = [
    tf.keras.callbacks.EarlyStopping(monitor='val_accuracy', 
                                     min_delta=0,
                                     patience=3,
                                     verbose=1,
                                     mode='auto',
                                     baseline=None,
                                     restore_best_weights=False,
                                     start_from_epoch=0),
    LRSchedulerCallback(lr_schedule)
]

# 编译并训练模型
classifier_model.compile(optimizer=optimizer, loss=loss, metrics=metrics)
history = classifier_model.fit(x=train_ds,
                               validation_data=val_ds,
                               epochs=epochs,
                               callbacks=callback,
                               steps_per_epoch=steps_per_epoch,
                               verbose=1)

方法2：延迟启用ReduceLROnPlateau（临时方案） #

如果想保留原有的optimization.create_optimizer，可以计算warmup所需的epoch数，让ReduceLROnPlateau在warmup结束后再生效：

# 计算warmup需要的epoch数
warmup_epochs = num_warmup_steps // steps_per_epoch + 1

# 修改ReduceLROnPlateau，设置start_from_epoch
callback = [
    tf.keras.callbacks.EarlyStopping(monitor='val_accuracy', 
                                     min_delta=0,
                                     patience=3,
                                     verbose=1,
                                     mode='auto',
                                     baseline=None,
                                     restore_best_weights=False,
                                     start_from_epoch=0),
    tf.keras.callbacks.ReduceLROnPlateau(
        monitor='val_loss', 
        factor=0.2,
        patience=5, 
        min_lr=0.001,
        start_from_epoch=warmup_epochs
    )
]

注意：此方法可能仍存在潜在问题，因为warmup结束后优化器的lr仍为WarmUp对象，回调修改时可能出现异常，仅作为临时过渡方案。

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