学习率调度器

DeepSpeed 提供了 LRRangeTest、OneCycle、WarmupLR、WarmupDecayLR、WarmupCosineLR 学习率调度器的实现。当使用 DeepSpeed 的学习率调度器（在 ds_config.json 文件中指定）时，DeepSpeed 会在每个训练步骤（当执行 model_engine.step() 时）调用调度器的 step() 方法。当不使用 DeepSpeed 的学习率调度器时

  • 如果调度计划应该在每个训练步骤执行，则用户可以在初始化 DeepSpeed 引擎时将调度器传递给 deepspeed.initialize，并让 DeepSpeed 管理其更新或保存/恢复。

  • 如果调度计划应该在任何其他间隔（例如，训练 epoch）执行，则用户不应在初始化期间将调度器传递给 DeepSpeed，并且必须显式地管理它。

LRRangeTest

class deepspeed.runtime.lr_schedules.LRRangeTest(optimizer: Optimizer, lr_range_test_min_lr: float = 0.001, lr_range_test_step_size: int = 2000, lr_range_test_step_rate: float = 1.0, lr_range_test_staircase: bool = False, last_batch_iteration: int = -1)[source]

根据学习率范围测试 (LRRT) 策略设置每个参数组的学习率。该策略从一个基值开始以恒定频率增加学习率，如论文 `A disciplined approach to neural network hyper-parameters: Part1`_ 中所述。

LRRT 策略用于查找在不发生发散的情况下训练模型的最大 LR，并且可用于配置循环 LR 调度的 LR 边界。

LRRT 在每个批次后更改学习率。step 应该在批次用于训练后调用。

参数

  • optimizer (Optimizer) – 包装的优化器。

  • lr_range_test_min_lr (float 或 list) – 初始学习率，它是每个参数组在范围测试中的下界。

  • lr_range_test_step_size (int) – 增加学习率的训练步骤间隔。默认值：2000

  • lr_range_test_step_rate (float) – 范围测试的缩放比率。默认值：1.0

  • lr_range_test_staircase (bool) – 以阶梯方式缩放，而不是连续缩放。默认值：False。

  • last_batch_iteration (int) – 最后一个批次的索引。此参数用于恢复训练作业。由于 step() 应该在每个批次后而不是每个 epoch 后调用，因此此数字表示计算的批次总数，而不是计算的 epoch 总数。当 last_batch_iteration=-1 时，调度从头开始。默认值：-1

示例

>>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
>>> scheduler = LRRangeTest(optimizer)
>>> data_loader = torch.utils.data.DataLoader(...)
>>> for epoch in range(10):
>>>     for batch in data_loader:
>>>         train_batch(...)
>>>         scheduler.step()

_神经网络超参数的纪律化方法：第 1 部分 - 学习率、批次大小、动量和权重衰减：https://arxiv.org/abs/1803.09820

OneCycle

class deepspeed.runtime.lr_schedules.OneCycle(optimizer, cycle_min_lr, cycle_max_lr, decay_lr_rate=0.0, cycle_first_step_size=2000, cycle_second_step_size=None, cycle_first_stair_count=0, cycle_second_stair_count=None, decay_step_size=0, cycle_momentum=True, cycle_min_mom=0.8, cycle_max_mom=0.9, decay_mom_rate=0.0, last_batch_iteration=-1)[source]

根据 1Cycle 学习率策略 (1CLR) 设置每个参数组的学习率。1CLR 是循环学习率 (CLR) 策略的一种变体，涉及一个循环，然后是衰减。该策略同时以恒定频率在两个边界之间循环学习率（和动量），如论文 神经网络超参数的纪律化方法 中所述。

1CLR 策略在每个批次后更改学习率。step 应该在批次用于训练后调用。

此实现改编自 github 存储库：`pytorch/pytorch`_

参数

  • optimizer (Optimizer) – 包装的优化器。

  • cycle_min_lr (float 或 list) – 初始学习率，它是每个参数组在循环中的下界。

  • cycle_max_lr (float 或 list) – 每个参数组在循环中的上限学习率。从功能上讲，它定义了循环幅度 (cycle_max_lr - cycle_min_lr)。任何循环中的 lr 都是 cycle_min_lr 和幅度的一些缩放的总和；因此，根据缩放函数，可能无法实际达到 cycle_max_lr。

  • decay_lr_rate (float) – 学习率的衰减率。默认值：0。

  • cycle_first_step_size (int) – 循环递增一半中的训练迭代次数。默认值：2000

  • cycle_second_step_size (int) – 循环递减一半中的训练迭代次数。如果 cycle_second_step_size 为 None，则将其设置为 cycle_first_step_size。默认值：None

  • cycle_first_stair_count (int) – 循环前半部分的阶梯数。这意味着

  • 0 (lr/mom 以阶梯方式更改。默认值) –

  • 禁用。 (表示阶梯) –

  • cycle_second_stair_count (int) – 循环后半部分的阶梯数。这意味着

  • 0 –

  • 禁用。 –

  • decay_step_size (int) – 在衰减阶段应用衰减的间隔。默认值：0，表示不衰减。

  • cycle_momentum (bool) – 如果 True，则动量与学习率反向循环，介于“cycle_min_mom”和“cycle_max_mom”之间。默认值：True

  • cycle_min_mom (float 或 list) – 初始动量，它是每个参数组在循环中的下界。默认值：0.8

  • cycle_max_mom (float 或 list) – 每个参数组在循环中的上限动量。从功能上讲，它定义了循环幅度 (cycle_max_mom - cycle_min_mom)。任何循环中的动量都是 cycle_max_mom 和幅度的一些缩放的差；因此，根据缩放函数，可能无法实际达到 cycle_min_mom。默认值：0.9

  • decay_mom_rate (float) – 动量的衰减率。默认值：0。

  • last_batch_iteration (int) – 最后一个批次的索引。此参数用于恢复训练作业。由于 step() 应该在每个批次后而不是每个 epoch 后调用，因此此数字表示计算的批次总数，而不是计算的 epoch 总数。当 last_batch_iteration=-1 时，调度从头开始。默认值：-1

示例

>>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
>>> scheduler = OneCycle(optimizer, 0.0001, 0.0010)
>>> data_loader = torch.utils.data.DataLoader(...)
>>> for epoch in range(10):
>>>     for batch in data_loader:
>>>         train_batch(...)
>>>         scheduler.step()

WarmupLR

class deepspeed.runtime.lr_schedules.WarmupLR(optimizer: Optimizer, warmup_min_lr: float = 0.0, warmup_max_lr: float = 0.001, warmup_num_steps: int = 1000, warmup_type: str = 'log', last_batch_iteration: int = -1)[source]

在 warmup_num_steps 步内，将每个参数组的学习率从最小学习率增加到最大学习率，然后固定在最大学习率。

参数

  • optimizer (Optimizer) – 包装的优化器。

  • warmup_min_lr (float 或 list) – 最小学习率。默认值：0

  • warmup_max_lr (float 或 list) – 最大学习率。默认值：0.001

  • warmup_num_steps (int) – 从 min_lr 到 max_lr 的预热步数。默认值：1000

  • {‘log’ (warmup_type) – 预热期间从 min_lr 到 max_lr 的递增函数。默认值：log

  • ‘linear’} – 预热期间从 min_lr 到 max_lr 的递增函数。默认值：log

  • last_batch_iteration (int) – 最后一个批次的索引。默认值：-1。

示例

>>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
>>> scheduler = WarmupLR(optimizer)
>>> data_loader = torch.utils.data.DataLoader(...)
>>> for epoch in range(10):
>>>     for batch in data_loader:
>>>         train_batch(...)
>>>         scheduler.step()

WarmupDecayLR

class deepspeed.runtime.lr_schedules.WarmupDecayLR(optimizer: Optimizer, total_num_steps: int, warmup_min_lr: float = 0.0, warmup_max_lr: float = 0.001, warmup_num_steps: int = 1000, warmup_type: str = 'log', last_batch_iteration: int = -1)[source]

在 warmup_num_steps 步内，将每个参数组的学习率从最小学习率增加到最大学习率，然后在剩余的训练步数中以线性速率衰减。

参数

  • optimizer (Optimizer) – 包装的优化器。

  • total_num_steps (int) – 总训练步数

  • warmup_min_lr (float 或 list) – 最小学习率。默认值：0

  • warmup_max_lr (float 或 list) – 最大学习率。默认值：0.001

  • warmup_num_steps (int) – 从 min_lr 到 max_lr 的预热步数。默认值：1000

  • {‘log’ (warmup_type) – 预热期间从 min_lr 到 max_lr 的递增函数。默认值：log

  • ‘linear’} – 预热期间从 min_lr 到 max_lr 的递增函数。默认值：log

  • last_batch_iteration (int) – 最后一个批次的索引。默认值：-1。

示例

>>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
>>> scheduler = WarmupDecayLR(optimizer, 1000000)
>>> data_loader = torch.utils.data.DataLoader(...)
>>> for epoch in range(10):
>>>     for batch in data_loader:
>>>         train_batch(...)
>>>         scheduler.step()

WarmupCosineLR

class deepspeed.runtime.lr_schedules.WarmupCosineLR(optimizer: Optimizer, total_num_steps: int, warmup_min_ratio: float = 0.0, warmup_num_steps: int = 1000, cos_min_ratio: float = 0.0001, warmup_type: str = 'log', last_batch_iteration: int = -1)[source]

在 warmup_num_steps 步内，将每个参数组的学习率从最小学习率比率增加到最大学习率比率，然后在剩余的训练步数中以余弦速率衰减到最小余弦比率。

参数

  • optimizer (Optimizer) – 包装的优化器。

  • total_num_steps (int) – 总训练步数

  • warmup_min_ratio (float 或 list) – 预热开始学习率比率。默认值：0

  • warmup_num_steps (int) – 从 warmup_min_ratio 预热到 1.0 的步数。默认值：1000

  • {‘log’ (warmup_type) – 预热期间从 min_lr 到 max_lr 的递增函数。默认值：log

  • ‘linear’} – 预热期间从 min_lr 到 max_lr 的递增函数。默认值：log

  • cos_min_ratio (float) – 余弦结束学习率比率。默认值：0.0001

  • last_batch_iteration (int) – 最后一个批次的索引。默认值：-1。

示例

>>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
>>> scheduler = WarmupCosineLR(optimizer, 1000000)
>>> data_loader = torch.utils.data.DataLoader(...)
>>> for epoch in range(10):
>>>     for batch in data_loader:
>>>         train_batch(...)
>>>         scheduler.step()