学习率范围测试

本教程展示了如何在 PyTorch 中进行学习率范围测试。

学习率范围测试 (LRRT)

学习率范围测试 (LRRT) 是一种发现可用于训练模型而不会发散的最大学习率值的方法。数据科学家通常对这些信息感兴趣，因为大的学习率比小的学习率能更快地加速模型收敛。此外，大的学习率在学习率调度器（如 CLR 和 1Cycle）中至关重要，这些调度器用于在大批量大小下有效地进行训练。DeepSpeed 为 PyTorch 框架中的模型训练提供了 LRRT 功能。

先决条件

要使用 DeepSpeed 的 LRRT 功能，您必须满足以下两个条件：

  1. 使用入门指南将 DeepSpeed 集成到您的训练脚本中。
  2. 将 LRRT 配置参数添加到模型的参数中。LRRT 参数定义如下。

LRRT 参数

LRRT 通过在预定义的间隔内，以预定义的量线性增加学习率。因此，LRRT 是一种学习率调度器，因为它定义了在模型训练期间学习率应如何以及何时改变。要配置 LRRT，您需要设置这些参数：

  1. lr_range_test_min_lr：训练的初始学习率 (float)
  2. lr_range_test_step_size：学习率的放大间隔，以训练步数定义 (integer)
  3. lr_range_test_step_rate：增加学习率的缩放因子 (float)
  4. lr_range_test_staircase：如果为 true，则学习率每 lr_range_test_step_size 个训练步改变一次；否则，学习率在每个训练步改变 (boolean)

所需的模型配置更改

我们将通过一个示例 LRRT 调度器来说明所需的模型配置更改，该调度器：

  1. 以 0.0001 的初始学习率开始训练
  2. 使用 5 的缩放率
  3. 使用 200 个训练步的缩放间隔
  4. 在每个训练步缩放学习率，即不使用阶梯式缩放

PyTorch

对于 PyTorch 模型，LRRT 作为学习率调度器实现，此功能在 PyTorch 1.0.1 及更高版本中可用。因此，您可以在模型配置中添加一个类型为 "LRRangeTest" 的 "scheduler" 条目，如下所示：

"scheduler": {
    "type": "LRRangeTest",
    "params": {
        "lr_range_test_min_lr": 0.0001,
        "lr_range_test_step_size": 200,
        "lr_range_test_step_rate": 5,
        "lr_range_test_staircase": false
    }
}

示例：为大批量大小调优

我们将通过我们调整一个内部生产模型以使其在更大的批量大小下高效收敛的经验片段，来展示 LRRT 如何使数据科学家受益，我们从一个 GPU（批量大小 512）扩展到四个 GPU（批量大小 2048）。我们的目标是使用相同的样本数量，以更大的批量大小训练模型，使其性能与更小的批量大小相匹配。这里的挑战是众所周知的大批量训练收敛缓慢问题。我们的方法是在 DeepSpeed 中使用 1Cycle 调度器来解决这个问题，并使用 LRRT 来配置该调度器。

在下面的图中，我们展示了如何使用 LRRT 发现用于批量大小为 2048 的有效训练的最大学习率。左图显示了在训练的前 9000 个批次中，大学习率对验证损失的影响。右图显示了在同一训练期间的学习率值。通过网格搜索，我们发现批量大小为 2048 的最佳固定学习率为 0.0002。蓝色线 (lr=0.0002) 表示使用此固定学习率进行训练。我们将两个 LRRT 调度器与此固定学习率进行比较。橙色线 (lr_range_test_step_rate=5) 和灰色线 (lr_range_test_step_rate=50) 表示使用相似的 LRRT 调度器进行训练，它们仅在 lr_range_test_step_rate 值上有所不同。尽管这些 LRRT 调度器从相同的基本学习率开始，但灰色线的学习率增长速度比橙色线快约 10 倍。此外，在所示数据点中，LRRT 调度器的学习率已经超过了蓝色线。我们随后将灰色线称为“快速增长”LRRT 调度器，橙色线称为“缓慢增长”LRRT 调度器。

validation_loss

从这个小例子中，我们得出以下观察结果。

  1. 更大的学习率显然能在一定程度上提升模型性能。快速增长的 LRRT 调度器在 3000 个批次后达到了 0.46 的验证损失，而固定学习率在 9000 个批次后仍未达到此水平。缓慢增长的 LRRT 直到 6000 个批次后才达到该分数，但它对固定学习率保持了不断增长的性能优势。

  2. 对于训练模型有用的学习率值存在一个上限。快速增长的 LRRT 调度器很快触及此边界并发生发散，而缓慢增长的 LRRT 也会因同样的原因在稍后发散。LRRT 帮助我们快速发现了这些边界，使用的训练数据不到 2%。这些边界对于构建学习率调度器是非常有用的信息。

LRRT 的这些观察结果帮助我们配置了 1Cycle 调度器的学习率边界和周期跨度，从而解决了问题，如下所示。

"OneCycle": {
    "cycle_min_lr": 0.002,
    "cycle_max_lr": 0.005,
    "cycle_first_step_size": 2000,
    "cycle_second_step_size": 2000,
    ...
}

根据我们的经验，这些是 1Cycle 调度器最关键的四个参数。

  1. 我们选择使用较慢的 LRRT 调度器 (lr_range_test_step_rate=5) 来设置 cycle_min_lr，因为它能达到最佳损失，而更快的调度器则很快发散。
  2. 我们将 cycle_min_lr 设置为 0.005，尽管图中显示在稍高的学习率下性能仍在提高。这是因为我们观察到，如果等到最大学习率，模型可能会达到发散点并且无法恢复。
  3. 由于学习率达到 0.005 需要 8000 个批次，我们将 cycle_first_step_size 和 (cycle_second_step_size) 设置为 2000，这是四个 GPU 处理 8000 个批次所需的步数。

我们希望这个简短的例子能激发您在应对自己独特的调优挑战时使用 LRRT 的灵感。

更新日期：