如何模拟实现ResNet论文中提及的学习率调整策略

规则说明 #

初始learning rate设为0.1，在迭代次数达到32k和48k时将其除以10，在64k迭代时终止训练，该训练终止节点是基于45k/5k的train/val拆分确定的。

对应调整逻辑示意图：

具体实现代码 #

1. 纯Python模拟实现 #

直接通过条件判断计算对应迭代次数的学习率：

def get_current_lr(iter_num: int) -> float:
    if iter_num >= 64000:
        raise ValueError("训练已达到终止迭代数64k")
    if iter_num < 32000:
        return 0.1
    elif 32000 <= iter_num < 48000:
        return 0.1 / 10
    else: # 48000 <= iter_num < 64000
        return 0.1 / 100

# 测试示例
test_iters = [10000, 32000, 40000, 48000, 60000]
for iter in test_iters:
    print(f"迭代次数{iter}对应的学习率为{get_current_lr(iter)}")

2. PyTorch框架内置实现 #

直接使用PyTorch提供的MultiStepLR调度器即可完美适配该策略：

import torch
from torch.optim.lr_scheduler import MultiStepLR

# 示例：定义一个简单的模型和优化器
model = torch.nn.Linear(10, 2)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

# 定义学习率调度器，milestones为调整节点，gamma为衰减系数
scheduler = MultiStepLR(optimizer, milestones=[32000, 48000], gamma=0.1)

# 训练循环示例
max_iter = 64000
for iter_num in range(max_iter):
    # 正常训练逻辑：前向传播、计算损失、反向传播、参数更新
    # optimizer.step()
    
    # 每次迭代后更新学习率
    scheduler.step()

3. TensorFlow/Keras实现 #

可以通过tf.keras.optimizers.schedules.PiecewiseConstantDecay实现：

import tensorflow as tf

boundaries = [32000, 48000]
values = [0.1, 0.01, 0.001]
lr_schedule = tf.keras.optimizers.schedules.PiecewiseConstantDecay(
    boundaries=boundaries,
    values=values
)

# 把调度器传入优化器即可
optimizer = tf.keras.optimizers.SGD(learning_rate=lr_schedule)

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