学习率调度器

DeepSpeed offers implementations of LRRangeTest, OneCycle, WarmupLR, WarmupDecayLR, WarmupCosineLR learning rate schedulers. When using a DeepSpeed’s learning rate scheduler (specified in the ds_config.json file), DeepSpeed calls the step() method of the scheduler at every training step (when model_engine.step() is executed). When not using a DeepSpeed’s learning rate scheduler:

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

  • 如果计划表应该在其它任何间隔执行（例如，训练周期），那么用户在初始化期间不应将调度器传递给DeepSpeed，并且必须明确管理它。

LR范围测试

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)[来源]

根据学习率范围测试（LRRT）策略设置每个参数组的学习率。该策略从基础值开始以恒定频率增加学习率，详细内容见论文A disciplined approach to neural network hyper-parameters: Part 1

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

LRRT 在每批次后改变学习率。 step 应在使用一个批次进行训练后调用。

Parameters

  • 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()

神经网络超参数的规范方法：第一部分 – 学习率、批量大小、动量和权重衰减： https://arxiv.org/abs/1803.09820

单周期

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)[来源]

根据1Cycle学习率策略（1CLR）设置每个参数组的学习率。1CLR是循环学习率（CLR）策略的一种变体，它涉及一个周期后跟随衰减。该策略以恒定频率在两个边界之间循环学习率（和动量），如论文A disciplined approach to neural network hyper-parameters中所述。

1CLR策略在每批处理后改变学习率。 step 应在使用一批数据进行训练后调用。

此实现改编自github仓库：PyTorch。

Parameters

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

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

  • cycle_max_lr (float 或 list) – 每个参数组在周期内的学习率上限。功能上，它定义了周期的振幅（cycle_max_lr - cycle_min_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 (学习率和动量以阶梯方式变化。默认) –

  • 禁用。 (意味着楼梯) –

  • 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()

预热学习率

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)[来源]

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

Parameters

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

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

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

  • warmup_num_steps (int) – 从最小学习率到最大学习率的预热步数。默认值：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()

预热衰减学习率

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)[来源]

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

Parameters

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

  • total_num_steps (int) – 训练步骤的总数

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

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

  • warmup_num_steps (int) – 从最小学习率到最大学习率的预热步数。默认值：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)[来源]

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

Parameters

  • 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()