YOLOv5训练时如何调整学习率？Scheduler使用指南

在深度学习中，学习率（Learning Rate）是一个非常重要的超参数，它决定了模型权重更新的速度。对于目标检测任务，尤其是使用YOLOv5这样的框架时，合理地调整学习率能够显著提升模型的训练效果和收敛速度。本文将深入探讨如何在YOLOv5训练过程中调整学习率，并介绍Scheduler的使用方法。

一、学习率的基本概念

学习率是梯度下降算法中的一个关键参数，用于控制每次迭代中模型参数更新的幅度。如果学习率过高，可能会导致训练不稳定或发散；如果学习率过低，则可能导致训练过程过于缓慢，甚至陷入局部最优解。

在YOLOv5中，默认的学习率策略已经经过优化，但根据具体任务需求，用户可能需要进一步调整学习率以获得更好的性能。

二、YOLOv5中的学习率调整机制

1. 默认学习率设置

YOLOv5默认采用了一种称为“余弦退火”（Cosine Annealing）的学习率调度器。这种调度器会随着训练轮次的增加逐渐降低学习率，从而帮助模型在后期更加精细地调整权重。

默认的学习率公式为： [ lr = lr_0 \times (1 + \cos(\frac{k}{n} \pi)) ] 其中：

• (lr_0) 是初始学习率；
• (k) 是当前的训练步数；
• (n) 是总训练步数。

2. 自定义学习率

如果默认的学习率策略不适用于特定任务，可以自定义学习率调度器。YOLOv5支持通过修改配置文件或代码来实现这一点。

三、Scheduler的使用指南

1. 修改hyp.scratch.yaml

YOLOv5中的超参数配置文件（如hyp.scratch.yaml）允许用户直接调整初始学习率。例如：

lr0: 0.01  # 初始学习率
lrf: 0.01  # 最终学习率的比例因子

• lr0：初始学习率。
• lrf：最终学习率的比例因子，最终学习率计算为lr_final = lr0 * lrf。

2. 使用自定义Scheduler

如果需要更复杂的调度器（如阶梯式衰减或指数衰减），可以通过修改train.py中的代码实现。以下是几种常见Scheduler的实现方式：

（1）阶梯式衰减（Step Decay）

阶梯式衰减会在指定的epoch点上将学习率按固定比例减少。例如：

from torch.optim.lr_scheduler import StepLR

# 定义优化器
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

# 定义Scheduler
scheduler = StepLR(optimizer, step_size=30, gamma=0.1)  # 每30个epoch将学习率乘以0.1

# 训练循环
for epoch in range(num_epochs):
    for batch in dataloader:
        optimizer.zero_grad()
        loss = model(batch)
        loss.backward()
        optimizer.step()
    scheduler.step()  # 更新学习率

（2）余弦退火（Cosine Annealing）

余弦退火是一种平滑的学习率下降策略，适用于大多数深度学习任务。YOLOv5默认使用该策略。如果需要手动实现，可以参考以下代码：

from torch.optim.lr_scheduler import CosineAnnealingLR

# 定义优化器
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)

# 定义Scheduler
scheduler = CosineAnnealingLR(optimizer, T_max=num_epochs, eta_min=0)

# 训练循环
for epoch in range(num_epochs):
    for batch in dataloader:
        optimizer.zero_grad()
        loss = model(batch)
        loss.backward()
        optimizer.step()
    scheduler.step()  # 更新学习率

（3）Warm-up策略

Warm-up策略通常用于大批次训练中，其核心思想是在训练初期逐步提高学习率，以避免模型因初始学习率过高而导致的训练不稳定。YOLOv5默认实现了Warm-up，可以通过修改hyp.scratch.yaml中的warmup_epochs参数来调整Warm-up的持续时间。

四、学习率调整的实际案例

1. 数据集较小的情况

当数据集较小时，模型容易过拟合。此时可以适当降低初始学习率，并延长Warm-up的时间。例如：

lr0: 0.001  # 降低初始学习率
warmup_epochs: 5  # 延长Warm-up时间

2. 数据集较大的情况

当数据集较大时，模型需要更快地收敛。此时可以适当提高初始学习率，并缩短Warm-up的时间。例如：

lr0: 0.01  # 提高初始学习率
warmup_epochs: 2  # 缩短Warm-up时间

五、可视化学习率变化

为了更好地观察学习率的变化，可以在训练过程中记录并绘制学习率曲线。以下是使用Matplotlib绘制学习率曲线的示例：

import matplotlib.pyplot as plt

# 假设我们有一个学习率列表
learning_rates = [0.01, 0.008, 0.006, 0.004, 0.002, 0.001]

plt.plot(range(len(learning_rates)), learning_rates, marker='o')
plt.title('Learning Rate Schedule')
plt.xlabel('Epochs')
plt.ylabel('Learning Rate')
plt.grid(True)
plt.show()

六、总结

通过本文，我们了解了YOLOv5中学习率调整的基本原理及Scheduler的使用方法。无论是默认的余弦退火策略还是自定义的阶梯式衰减或Warm-up策略，合理的选择和调整学习率都能显著提升模型的训练效果。