1. 为什么我们需要Warmup与Cosine Annealing？

如果你在训练深度学习模型时，感觉模型一开始就“跑偏”了，或者训练损失下降得特别慢，那很可能就是学习率没调好。我刚开始炼丹那会儿，也总在这个问题上栽跟头。后来我发现，很多高手都在用“预热（Warmup）+余弦退火（Cosine Annealing）”这套组合拳，它能让模型训练得更稳、效果更好。

简单来说，Warmup就像运动前的热身。模型刚初始化时，它的参数是随机的，梯度可能非常大。如果一上来就用很大的学习率去更新，步子迈得太大，很容易导致训练不稳定，甚至直接“跑飞”。Warmup就是在训练初期，让学习率从一个很小的值（比如0）开始，线性地（或其他方式）逐渐增加到你设定的初始学习率。这个过程给了模型一个“缓冲期”，让它先适应一下数据，把参数分布调整得平稳一些，后续再用大学习率更新时，就不容易出问题了。

而Cosine Annealing则像是给训练过程安排了一个“优雅的谢幕”。传统的学习率衰减，比如每隔固定轮次减半，是阶梯式下降的，这种突变有时会让模型不适应。余弦退火则不同，它让学习率按照余弦函数的曲线，从最大值平滑地下降到最小值。这个过程非常柔和，能让模型在训练后期，慢慢地、稳定地收敛到一个更优的解，而不是在最小值点附近来回震荡。

PyTorch官方提供了CosineAnnealingLR和CosineAnnealingWarmRestarts等调度器，但它们都没有内置Warmup阶段。想把两者结合起来，最灵活、最强大的工具就是LambdaLR。它允许你通过一个自定义的lambda函数，完全掌控学习率变化的每一个细节。接下来，我就带你彻底搞懂怎么用LambdaLR把这两者无缝衔接起来。

2. 深入理解LambdaLR：你的学习率画笔

torch.optim.lr_scheduler.LambdaLR是PyTorch学习率调度器家族里最“自由”的一个。它的核心参数是lr_lambda，可以是一个函数，也可以是一个函数列表。这个函数的作用是：在每一次调度器更新时，计算一个乘数因子，然后用这个因子去乘以优化器中每个参数组（param_group）的初始学习率（base_lr）。

这里有个关键点必须理解：LambdaLR修改的是学习率乘数，而不是直接设置学习率绝对值。公式是：当前学习率 = base_lr * lr_lambda(current_step)。很多新手在这里踩坑，写lambda函数时直接返回想要的学习率数值，结果发现实际学习率不对，就是因为忘了它是个乘数。

LambdaLR的强大之处在于它能针对不同的参数组设置不同的调整策略。在定义优化器时，我们可以把模型的参数分成多个组，每个组可以有自己的初始学习率和其他超参数。然后，给lr_lambda传递一个函数列表，列表中的每个函数会依次应用到对应的参数组上。这个功能在微调模型时特别有用，比如我们希望骨干网络（backbone）的学习率小一点、慢一点衰减，而新加的分类头（head）学习率大一点、衰减快一点。

下面我们通过一个最简单的例子，看看LambdaLR是怎么工作的：

import torch
import torch.nn as nn

# 定义一个简单的模型
model = nn.Linear(10, 2)

# 创建优化器，并设置两个参数组，它们的基础学习率不同
optimizer = torch.optim.SGD([
    {'params': model.weight, 'lr': 0.1},   # 参数组0，base_lr=0.1
    {'params': model.bias, 'lr': 0.01}     # 参数组1，base_lr=0.01
])

# 定义两个lambda函数作为乘数
# 参数组0：学习率每步变为之前的0.9倍
lambda0 = lambda epoch: 0.9 ** epoch
# 参数组1：学习率每步线性减少到0
lambda1 = lambda epoch: max(1.0 - epoch / 10.0, 0)

# 创建LambdaLR调度器
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=[lambda0, lambda1])

# 模拟训练，打印学习率变化
for epoch in range(5):
    print(f'Epoch {epoch}:')
    for i, group in enumerate(optimizer.param_groups):
        print(f'  Param Group {i} lr: {group["lr"]:.6f}')
    scheduler.step()

运行这段代码，你会看到两个参数组的学习率按照我们定义的规则独立变化。理解了LambdaLR这个“乘数”本质和分组控制能力，我们就能用它来绘制任意形状的学习率曲线了。

3. 动手实现：从零编写Warmup + Cosine Annealing函数

现在，我们来把Warmup和Cosine Annealing的数学公式，翻译成LambdaLR能听懂的lambda函数。首先，我们明确一下整个策略的过程：

1. Warmup阶段（前 warmup_iter 步）：学习率从0线性增长到最大学习率 lr_max。

   • 乘数公式：lambda(current_iter) = current_iter / warmup_iter

2. Cosine Annealing阶段（warmup_iter 步之后）：学习率从 lr_max 按照余弦曲线平滑下降到最小学习率 lr_min。

   • 余弦退火公式：lr_min + 0.5 * (lr_max - lr_min) * (1 + cos(π * progress))
   • 其中 progress = (current_iter - warmup_iter) / (T_max - warmup_iter)，代表余弦阶段的进度。

这里有一个至关重要的细节：我们的lambda函数返回的是乘数。在Warmup阶段，base_lr * (current_iter/warmup_iter) 确实能从0增长到base_lr。但是进入Cosine阶段后，公式计算出来的是绝对的学习率值（比如0.05），而不是乘数。因此，我们需要将这个绝对学习率值，除以该参数组的base_lr，来得到正确的乘数。

假设我们的base_lr（即lr_max）是0.1，在某个Cosine阶段计算出的学习率是0.05，那么乘数应该是 0.05 / 0.1 = 0.5。这样，LambdaLR执行 0.1 * 0.5 才能得到正确的0.05。

结合以上两点，我们可以写出完整的lambda函数：

import math

def get_cosine_annealing_with_warmup_lambda(warmup_iter, T_max, lr_min, base_lr):
    """
    返回一个可用于LambdaLR的lambda函数，实现Warmup + Cosine Annealing。
    参数:
        warmup_iter: Warmup的步数（迭代次数）。
        T_max: 总的学习率周期（通常是总迭代步数）。注意，这里T_max是包含warmup步数的总步数。
        lr_min: 余弦退火下降到的最小学习率。
        base_lr: 该参数组的基础学习率（即lr_max）。
    """
    def lr_lambda(current_iter):
        # Warmup阶段
        if current_iter < warmup_iter:
            # 线性增长乘数：从0到1
            return float(current_iter) / float(max(1, warmup_iter))
        # Cosine Annealing阶段
        else:
            # 计算余弦阶段的进度，确保不超过1
            progress = float(current_iter - warmup_iter) / float(max(1, T_max - warmup_iter))
            progress = min(1.0, progress)  # 防止progress超过1
            # 计算余弦退火后的绝对学习率
            cosine_lr = lr_min + 0.5 * (base_lr - lr_min) * (1.0 + math.cos(math.pi * progress))
            # 关键：将绝对学习率转换为乘数
            return cosine_lr / base_lr
    return lr_lambda

这个函数封装了所有逻辑，使用起来非常清晰。我们把它应用到之前的例子中，并加上可视化，看看学习率曲线到底长什么样。

4. 完整代码示例与可视化分析

理论说再多，不如跑段代码看看。下面我给出一个完整的、可以直接运行的例子，它包含了模型定义、优化器设置、我们的组合调度器，以及用Matplotlib绘制学习率变化曲线的部分。

import torch
import torch.nn as nn
import math
import matplotlib.pyplot as plt

# 1. 使用上面的函数定义
def get_cosine_annealing_with_warmup_lambda(warmup_iter, T_max, lr_min, base_lr):
    def lr_lambda(current_iter):
        if current_iter < warmup_iter:
            return float(current_iter) / float(max(1, warmup_iter))
        else:
            progress = float(current_iter - warmup_iter) / float(max(1, T_max - warmup_iter))
            progress = min(1.0, progress)
            cosine_lr = lr_min + 0.5 * (base_lr - lr_min) * (1.0 + math.cos(math.pi * progress))
            return cosine_lr / base_lr
    return lr_lambda

# 2. 创建一个模拟模型和优化器
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(784, 256)
        self.fc2 = nn.Linear(256, 10)
    def forward(self, x):
        x = torch.relu(self.fc1(x))
        return self.fc2(x)

model = SimpleModel()

# 3. 设置优化器，并创建两个参数组（模拟不同层不同学习率）
# 假设fc1层我们想用较小的基础学习率，且不进行warmup（仅cosine衰减）
# fc2层用较大的基础学习率，并进行完整的warmup+cosine
optimizer = torch.optim.AdamW([
    {'params': model.fc1.parameters(), 'lr': 1e-4},  # 参数组0
    {'params': model.fc2.parameters(), 'lr': 1e-3},  # 参数组1
])

# 4. 定义调度器参数
total_iterations = 200  # 总训练迭代次数（假设）
warmup_iterations = 20   # Warmup迭代次数
min_lr = 1e-6           # 余弦退火下降到的最小学习率

# 5. 为每个参数组创建对应的lambda函数
# 参数组0: 只有Cosine Annealing (从1e-4降到min_lr)，没有Warmup。我们可以将warmup_iter设为0。
lambda0 = get_cosine_annealing_with_warmup_lambda(
    warmup_iter=0,  # 无warmup
    T_max=total_iterations,
    lr_min=min_lr,
    base_lr=1e-4
)
# 参数组1: 完整的Warmup + Cosine Annealing (从1e-3降到min_lr)
lambda1 = get_cosine_annealing_with_warmup_lambda(
    warmup_iter=warmup_iterations,
    T_max=total_iterations,
    lr_min=min_lr,
    base_lr=1e-3
)

# 6. 创建LambdaLR调度器
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=[lambda0, lambda1])

# 7. 模拟训练过程，记录学习率
lr_history_0 = []
lr_history_1 = []
for step in range(total_iterations):
    # 记录当前步的学习率
    lr_history_0.append(optimizer.param_groups[0]['lr'])
    lr_history_1.append(optimizer.param_groups[1]['lr'])
    # 模拟训练步骤...
    # optimizer.zero_grad()
    # loss.backward()
    # optimizer.step()
    # 更新学习率调度器（按迭代次数更新）
    scheduler.step()

# 8. 可视化
plt.figure(figsize=(10, 6))
plt.plot(range(total_iterations), lr_history_0, label='Param Group 0 (FC1, lr=1e-4, No Warmup)')
plt.plot(range(total_iterations), lr_history_1, label='Param Group 1 (FC2, lr=1e-3, With Warmup)', linestyle='--')
plt.axvline(x=warmup_iterations, color='gray', linestyle=':', label='Warmup End')
plt.xlabel('Training Iteration')
plt.ylabel('Learning Rate')
plt.title('Warmup + Cosine Annealing LR Schedule (Using LambdaLR)')
plt.legend()
plt.grid(True, which='both', linestyle='--', alpha=0.5)
plt.yscale('log')  # 使用对数坐标更清晰地观察小学习率变化
plt.show()

运行这段代码，你会得到一张清晰的学习率变化图。从图中你可以明显看到：

• 参数组1（虚线）的学习率从0开始，在20次迭代内线性增长到1e-3，然后平滑地沿余弦曲线下降。
• 参数组0（实线）因为没有设置Warmup，直接从1e-4开始沿余弦曲线下降。
• 两条曲线最终都收敛到我们设定的最小值1e-6附近。

这种可视化能帮你直观地确认调度器是否按预期工作，对于调试超参数（比如warmup_iterations和T_max的比例）非常有帮助。

5. 高级技巧与实战避坑指南

在实际项目中，用好LambdaLR还需要注意一些细节和技巧，这些往往是在踩过几次坑之后才总结出来的。

技巧一：step()的调用时机 LambdaLR默认是按epoch来更新的（即step()函数通常在每个epoch结束后调用）。但在我们上面的实现中，lambda函数接收的参数current_iter实际上是内部记录的一个计数器，默认每调用一次step()就加1。如果你想实现按迭代次数（iteration/batch）更新学习率，就需要在每个batch处理后调用scheduler.step()。这时，你需要确保total_iterations等参数是基于总batch数来设置的，而不是总epoch数。这是新手常混淆的地方。

技巧二：与CosineAnnealingWarmRestarts结合 有时我们不仅想要一次余弦下降，还希望有“热重启”（Warm Restarts），即在学习率下降到最小值后，突然跳回一个较高的值重新开始下降，这有助于模型跳出局部最优。PyTorch提供了CosineAnnealingWarmRestarts调度器，但它同样没有内置Warmup。我们可以用LambdaLR实现初始的Warmup，然后切换到官方的CosineAnnealingWarmRestarts。更优雅的方式是使用PyTorch的SequentialLR，它可以将多个调度器按顺序串联。下面是一个概念示例：

import torch.optim.lr_scheduler as lr_scheduler

# 假设我们已经有了optimizer
warmup_scheduler = LambdaLR(optimizer, lr_lambda=lambda step: step / warmup_steps if step < warmup_steps else 1.0)
cosine_restart_scheduler = lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0=restart_cycle_length)

# 使用SequentialLR串联：先执行warmup_scheduler（持续warmup_steps步），再执行cosine_restart_scheduler
from torch.optim.lr_scheduler import SequentialLR
combined_scheduler = SequentialLR(
    optimizer,
    schedulers=[warmup_scheduler, cosine_restart_scheduler],
    milestones=[warmup_steps]  # 在warmup_steps步后切换到第二个调度器
)

避坑指南：参数组与lambda函数的匹配 务必确保lr_lambda列表的长度与优化器中参数组（param_groups）的数量一致。如果只有一个lambda函数，它会应用到所有参数组。如果参数组数量为3，而lr_lambda列表只有2个函数，程序会报错。我建议在创建调度器后，立即打印一下len(optimizer.param_groups)和len(scheduler.lr_lambdas)来做个检查。

避坑指南：恢复训练时的状态 当你要从检查点（checkpoint）恢复训练时，除了加载模型参数和优化器状态，千万别忘了加载调度器的状态。调度器内部有last_epoch等状态变量，如果不加载，学习率调度会从头开始，打乱原有的计划。保存和加载的代码如下：

# 保存检查点
checkpoint = {
    'model_state_dict': model.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'scheduler_state_dict': scheduler.state_dict(),
    'epoch': epoch,
    # ... 其他信息
}
torch.save(checkpoint, 'checkpoint.pth')

# 加载检查点
checkpoint = torch.load('checkpoint.pth')
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
# 接下来从 checkpoint['epoch'] 开始继续训练

6. 在不同场景下的参数调优建议

Warmup和Cosine Annealing虽好，但里面的超参数（warmup_iterations， T_max， lr_min）怎么设呢？这没有绝对的金标准，但有一些经验法则可以参考。

• Warmup迭代次数（warmup_iterations）：通常设置为总迭代次数的5%到10%。例如，如果你计划训练100个epoch，每个epoch有1000个batch，那么总迭代就是10万次，Warmup可以设为5000到10000次迭代。对于非常大的模型（如Transformer）或非常不稳定的任务，可以适当延长Warmup。你可以先设一个值，通过上面的可视化看看学习率曲线，感觉一下“热身”阶段是否足够平缓。

• 余弦周期（T_max）：在我们的实现中，T_max通常就是总训练迭代次数。这意味着整个训练过程就是一个完整的“预热-余弦下降”大周期。这也是最常见、最稳定的用法。如果你想尝试带重启的余弦退火，那么T_max就是每个周期的长度。

• 最小学习率（lr_min）：这个值通常设为基础学习率（base_lr）的1/10到1/1000。例如，base_lr是1e-3，lr_min可以设为1e-4或1e-6。设得太小（比如0），训练后期可能完全停止更新；设得太大，又起不到让模型充分收敛的效果。一个稳妥的做法是参考你领域内顶尖论文或开源代码的设置。

• 参数组差异化配置：这是LambdaLR发挥威力的地方。比如在图像分类任务中微调预训练模型，我通常会这样设置：

  • 骨干网络：base_lr较小（如1e-5），warmup_iterations较短甚至为0（因为预训练权重已经比较稳定），衰减可以慢一点（比如lr_min设为base_lr的1/10）。
  • 新添加的分类层：base_lr较大（如1e-3），进行完整的、时间较长的Warmup，然后衰减到更小的lr_min（如1e-6）。 通过这种差异化的精细控制，往往能比所有层“一刀切”获得更好的微调效果。

最后，记住一点：所有超参数的建议都是起点。最可靠的方法还是在你的验证集上，围绕这些默认值进行小范围的网格搜索或随机搜索，找到最适合你具体任务、具体数据的那组“神奇数字”。毕竟，调参本身就是深度学习工程师最重要的手艺之一。