学习率范围测试

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

学习率范围测试 (LRRT)

学习率范围测试(LRRT)是一种用于发现可用于训练模型而不导致发散的最大学习率值的方法。数据科学家通常对此信息感兴趣,因为较大的学习率比小的学习率能更快地使模型收敛。此外,较大的学习率在学习率调度中至关重要,例如CLR1Cycle,这些调度用于在大批量大小下有效训练。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: 如果为真,学习率将在每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)线表示使用仅在lr_range_test_step_rate值上不同的类似LRRT计划进行训练。尽管LRRT计划从相同的基础学习率开始,但灰色线的学习率增长速度比橙色线快约10倍。此外,在所呈现的数据点中,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. 由于需要8000批次才能使学习率降至0.005,我们将cycle_first_step_sizecycle_second_step_size设置为2000,这是四个GPU处理8000批次所需的步数。

我们希望这个简短的示例能激发您使用LRRT解决自己独特调优挑战的想象力。

更新: