如何获取MultiOptimizer中AdamW对应不同层的学习率并可视化

当使用tfa.optimizers.MultiOptimizer时，主优化器没有统一的lr属性，你需要从optimizer_specs中提取每个子优化器的学习率。每个optimizer_specs元素是一个字典，其中optimizer字段对应你定义的AdamW实例，直接访问这个实例的learning_rate即可获取对应层的学习率。

修改你的自定义TensorBoard回调如下：

from keras import backend as K
from keras.callbacks import TensorBoard

class LRTensorBoard(TensorBoard):
    def __init__(self, log_dir, **kwargs):
        super().__init__(log_dir=log_dir, **kwargs)

    def on_epoch_end(self, epoch, logs=None):
        logs = logs or {}
        # 提取预训练层对应的学习率
        pretrained_lr = K.eval(self.model.optimizer.optimizer_specs[0]['optimizer'].learning_rate)
        # 提取自定义层对应的学习率
        custom_lr = K.eval(self.model.optimizer.optimizer_specs[1]['optimizer'].learning_rate)
        # 将两个学习率写入logs，TensorBoard会分别显示
        logs.update({'lr_pretrained': pretrained_lr, 'lr_custom': custom_lr})
        super().on_epoch_end(epoch, logs)

然后在训练时正常使用这个回调：

model.fit(..., callbacks=[LRTensorBoard(path)])

关键说明 #

• optimizer_specs的顺序和你创建optimizers_and_layers时的顺序一致：第一个元素对应预训练层的AdamW优化器，第二个对应自定义层的。
• 如果你的学习率是动态变化的（比如使用学习率调度器），K.eval()会自动获取当前epoch的实际学习率值。
• 写入logs时使用不同的键名（比如lr_pretrained和lr_custom），TensorBoard会在标量面板中显示两条独立的学习率曲线，方便对比。

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