学习率调度器

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）策略设置每个参数组的学习率。该策略以固定频率从基值开始增加学习率，具体细节参见论文神经网络超参数的规范方法：第一部分

LRRT 策略用于找到训练模型而不发散的最大学习率（LR），并可用于配置循环学习率（Cyclic LR）调度器的 LR 边界。

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

参数

optimizer (Optimizer) – 包装的优化器。
lr_range_test_min_lr (浮点数 或列表) – 初始学习率，它是每个参数组在范围测试中的下限。
lr_range_test_step_size (整数) – 增加学习率的训练步数间隔。默认值：2000
lr_range_test_step_rate (浮点数) – 范围测试的缩放率。默认值：1.0
lr_range_test_staircase (布尔值) – 阶梯式缩放，而非连续缩放。默认值：False。
last_batch_iteration (整数) – 最后一个批次的索引。此参数在恢复训练任务时使用。由于 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()

神经网络超参数的规范方法：第一部分——学习率、批大小、动量和权重衰减：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。

参数

optimizer (Optimizer) – 包装的优化器。
cycle_min_lr (浮点数 或列表) – 初始学习率，它是每个参数组在周期中的下限。
cycle_max_lr (浮点数 或列表) – 每个参数组在周期中的学习率上限。功能上，它定义了周期幅度（`cycle_max_lr - cycle_min_lr`）。任何周期中的学习率是 `cycle_min_lr` 与幅度某种缩放的和；因此，根据缩放函数，`cycle_max_lr` 可能实际上无法达到。
decay_lr_rate (浮点数) – 学习率的衰减率。默认值：0。
cycle_first_step_size (整数) – 周期上升阶段的训练迭代次数。默认值：2000
cycle_second_step_size (整数) – 周期下降阶段的训练迭代次数。如果 `cycle_second_step_size` 为 None，则将其设置为 `cycle_first_step_size`。默认值：None
cycle_first_stair_count (整数) – 周期第一阶段的阶梯数。这意味着
0 (学习率/动量以阶梯式方式改变。默认值) –
禁用。 (表示阶梯式) –
cycle_second_stair_count (整数) – 周期第二阶段的阶梯数。这意味着
0 –
禁用。 –
decay_step_size (整数) – 衰减阶段应用衰减的间隔。默认值：0，表示不衰减。
cycle_momentum (布尔值) – 如果为 True，动量会与学习率反向循环，在 'cycle_min_mom' 和 'cycle_max_mom' 之间。默认值：True
cycle_min_mom (浮点数 或列表) – 初始动量，它是每个参数组在周期中的下限。默认值：0.8
cycle_max_mom (浮点数 或列表) – 每个参数组在周期中的动量上限。功能上，它定义了周期幅度（`cycle_max_mom - cycle_min_mom`）。任何周期中的动量是 `cycle_max_mom` 与幅度某种缩放的差值；因此，根据缩放函数，`cycle_min_mom` 可能实际上无法达到。默认值：0.9
decay_mom_rate (浮点数) – 动量的衰减率。默认值：0。
last_batch_iteration (整数) – 最后一个批次的索引。此参数在恢复训练任务时使用。由于 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 (浮点数 或列表) – 最小学习率。默认值：0
warmup_max_lr (浮点数 或列表) – 最大学习率。默认值：0.001
warmup_num_steps (整数) – 从 `min_lr` 预热到 `max_lr` 的步数。默认值：1000
{‘log’ (warmup_type) – 预热期间从 `min_lr` 到 `max_lr` 的递增函数。默认值：log
‘linear’} – 预热期间从 `min_lr` 到 `max_lr` 的递增函数。默认值：log
last_batch_iteration (整数) – 最后一个批次的索引。默认值：-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 (整数) – 总训练步数
warmup_min_lr (浮点数 或列表) – 最小学习率。默认值：0
warmup_max_lr (浮点数 或列表) – 最大学习率。默认值：0.001
warmup_num_steps (整数) – 从 `min_lr` 预热到 `max_lr` 的步数。默认值：1000
{‘log’ (warmup_type) – 预热期间从 `min_lr` 到 `max_lr` 的递增函数。默认值：log
‘linear’} – 预热期间从 `min_lr` 到 `max_lr` 的递增函数。默认值：log
last_batch_iteration (整数) – 最后一个批次的索引。默认值：-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 (整数) – 总训练步数
warmup_min_ratio (浮点数 或列表) – 预热起始学习率比率。默认值：0
warmup_num_steps (整数) – 从 `warmup_min_ratio` 预热到 1.0 的步数。默认值：1000
{‘log’ (warmup_type) – 预热期间从 `min_lr` 到 `max_lr` 的递增函数。默认值：log
‘linear’} – 预热期间从 `min_lr` 到 `max_lr` 的递增函数。默认值：log
cos_min_ratio (浮点数) – 余弦结束学习率比率。默认值：0.0001
last_batch_iteration (整数) – 最后一个批次的索引。默认值：-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()