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

引发：

RuntimeError – 如果学习率调度器尚未设置。

get_optim_key(key: str) → Any

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

load_state_dict(optim_ckpt_map: Dict[str, Any])

从此类的 state_dict 方法产生的状态字典中加载优化器状态。

参数：

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

引发：

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, 可选) – 当前的 epoch 编号。默认为 None。

引发：

RuntimeError – 如果学习率调度器尚未设置。