学习率调度器

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

如果调度应该在每次训练步骤执行，那么用户可以在初始化 DeepSpeed 引擎时将调度器传递给 deepspeed.initialize，并让 DeepSpeed 管理它以进行更新或保存/恢复。
如果调度应该在任何其他间隔执行（例如，训练时期），那么用户不应在初始化期间将调度器传递给 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 调度的 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() 应该在每个批次后而不是在每个时期后调用，因此此数字表示计算的总批次数，而不是计算的总时期数。当 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()

_A disciplined approach to neural network hyper-parameters: Part 1 – learning rate, batch size, momentum, and weight decay: 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) 策略的一种变体，它涉及一个循环，然后是衰减。该策略以恒定频率同时在两个边界之间循环学习率（和动量），如论文 A disciplined approach to neural network hyper-parameters 中所述。

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() 应该在每个批次后而不是在每个时期后调用，因此此数字表示计算的总批次数，而不是计算的总时期数。当 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 个步骤内将每个参数组的学习率从最小 lr 增加到最大 lr，然后固定在最大 lr 上。

参数

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 步内，将每个参数组的学习率从最小 lr 增加到最大 lr，然后在剩余的训练步骤中以线性速率衰减。

参数

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 步内，将每个参数组的学习率从最小 lr 比例增加到最大 lr 比例，然后在剩余的训练步骤中以余弦速率衰减到最小余弦比例。

参数

optimizer (Optimizer) – 包装的优化器。
total_num_steps (int) – 训练步骤的总数
warmup_min_ratio (float or 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()