OptimizerInBackwardWrapper

class torchtune.training.OptimizerInBackwardWrapper(optim_map: Dict[str, Optimizer])

一个精简的类，旨在为在反向传播中运行的优化器进行检查点保存和加载。用法仅限于以下情况

注意

此包装器仅适用于单设备用例。不支持分布式用例，例如 FSDP，它需要专门的优化器状态检查点。

参数:

optim_map (Dict[str, torch.optim.Optimizer]) – 从参数名称到优化器的映射。

示例

>>> optim_dict = {
>>>     p: config.instantiate(cfg_optimizer, [p])
>>>     for p in self._model.parameters()
>>> }
>>>
>>> # Save checkpoint
>>> ckpt = OptimizerInBackwardWrapper(optim_dict).state_dict()
>>> torch.save("/tmp/optim_ckpt", ckpt)
>>>
>>> # Load checkpoint
>>> placeholder_optim_dict = {
>>>     p: config.instantiate(cfg_optimizer, [p])
>>>     for p in self._model.parameters()
>>> }
>>>
>>> wrapper = OptimInBackwardWrapper(placeholder_optim_dict)
>>>
>>> # load_state_dict expects a dict produced by this class's
>>> # state_dict method.
>>> wrapper.load_state_dict(torch.load("/tmp/optim_ckpt"))
>>> # placeholder_optim_dict now has updated optimizer states.

get_last_lr() → float

如果存在调度器，则获取调度器中的最后一个学习率。

返回:

最后一个学习率。

返回类型:

float

Raises:

RuntimeError – 如果未设置 LR 调度器。

get_optim_key(key: str) → Any

从在反向传播中运行的任意优化器返回键的值。请注意，这假定反向传播中的所有优化器对于该键都具有相同的值，即使用相同的超参数初始化。

load_state_dict(optim_ckpt_map: Dict[str, Any])

从该类的 state_dict 方法生成的状态字典加载优化器状态。

参数:

optim_ckpt_map (Dict[str, Any]) – 状态字典，将参数名称映射到优化器状态。

Raises:

RuntimeError – 如果优化器状态字典不包含所有预期的参数。

set_lr_scheduler(lr_scheduler: LRScheduler) → None

设置学习率调度器并修改其 step 方法以更新所有优化器。

参数:

lr_scheduler (LRScheduler) – 要使用的学习率调度器。

state_dict() → Dict[str, Any]

返回状态字典，将参数名称映射到优化器状态。此 state_dict 只能由此类加载。

返回:

状态字典，将参数名称映射到优化器状态。

返回类型:

Dict[str, Any]

step_lr_scheduler(epoch: int = None)

如果存在学习率调度器，则执行步进操作。

参数:

epoch (int, optional) – 当前 epoch 编号。默认为 None。

Raises:

RuntimeError – 如果未设置 LR 调度器。