Adam

classtorch.optim.Adam(params, lr=0.001, betas=(0.9, 0.999), eps=1e-08, weight_decay=0, amsgrad=False, *, foreach=None, maximize=False, capturable=False, differentiable=False, fused=None)[源代码]

实现Adam算法。

$\begin{aligned} &\rule{110mm}{0.4pt} \\ &\textbf{input} : \gamma \text{ (lr)}, \beta_1, \beta_2 \text{ (betas)},\theta_0 \text{ (params)},f(\theta) \text{ (objective)} \\ &\hspace{13mm} \lambda \text{ (weight decay)}, \: \textit{amsgrad}, \:\textit{maximize} \\ &\textbf{initialize} : m_0 \leftarrow 0 \text{ ( first moment)}, v_0\leftarrow 0 \text{ (second moment)},\: \widehat{v_0}^{max}\leftarrow 0\\[-1.ex] &\rule{110mm}{0.4pt} \\ &\textbf{for} \: t=1 \: \textbf{to} \: \ldots \: \textbf{do} \\ &\hspace{5mm}\textbf{if} \: \textit{maximize}: \\ &\hspace{10mm}g_t \leftarrow -\nabla_{\theta} f_t (\theta_{t-1}) \\ &\hspace{5mm}\textbf{else} \\ &\hspace{10mm}g_t \leftarrow \nabla_{\theta} f_t (\theta_{t-1}) \\ &\hspace{5mm}\textbf{if} \: \lambda \neq 0 \\ &\hspace{10mm} g_t \leftarrow g_t + \lambda \theta_{t-1} \\ &\hspace{5mm}m_t \leftarrow \beta_1 m_{t-1} + (1 - \beta_1) g_t \\ &\hspace{5mm}v_t \leftarrow \beta_2 v_{t-1} + (1-\beta_2) g^2_t \\ &\hspace{5mm}\widehat{m_t} \leftarrow m_t/\big(1-\beta_1^t \big) \\ &\hspace{5mm}\widehat{v_t} \leftarrow v_t/\big(1-\beta_2^t \big) \\ &\hspace{5mm}\textbf{if} \: amsgrad \\ &\hspace{10mm}\widehat{v_t}^{max} \leftarrow \mathrm{max}(\widehat{v_t}^{max}, \widehat{v_t}) \\ &\hspace{10mm}\theta_t \leftarrow \theta_{t-1} - \gamma \widehat{m_t}/ \big(\sqrt{\widehat{v_t}^{max}} + \epsilon \big) \\ &\hspace{5mm}\textbf{else} \\ &\hspace{10mm}\theta_t \leftarrow \theta_{t-1} - \gamma \widehat{m_t}/ \big(\sqrt{\widehat{v_t}} + \epsilon \big) \\ &\rule{110mm}{0.4pt} \\[-1.ex] &\bf{return} \: \theta_t \\[-1.ex] &\rule{110mm}{0.4pt} \\[-1.ex] \end{aligned}$

有关算法的更多细节，请参阅 《Adam: 随机优化方法》。

参数

• params (可迭代对象) – 需要优化的参数或定义参数组的字典的可迭代对象

• lr (float, Tensor, optional) – 学习率（默认值：1e-3）。目前所有实现还不支持使用张量形式的 LR。如果没有指定 fused=True 或 capturable=True，请使用 float 类型的 LR。

• betas (Tuple[float, float], optional) – 用于计算梯度及其平方的滑动平均系数（默认值：(0.9, 0.999)）

• eps (float, 可选) – 用于提高数值稳定性的添加到分母中的项（默认值：1e-8）

• weight_decay (float, optional) – L2惩罚（权重衰减）（默认值：0）

• amsgrad (bool, 可选) – 是否使用论文《On the Convergence of Adam and Beyond》中提出的 AMSGrad 算法变体（默认值：False）

• foreach (bool, 可选) – 是否使用优化器的 foreach 实现。如果用户未指定（即 foreach 为 None），我们将尝试在 CUDA 上使用 foreach 替代 for 循环实现，因为前者通常性能更好。需要注意的是，foreach 实现比单个张量版本多占用约 sizeof(params) 的峰值内存，这是因为中间结果是 tensorlist 而不是单一的张量。如果内存受限，请减少每次通过优化器批量处理的参数数量或将此标志设置为 False（默认值：None）

• maximize (bool, optional) – 是否相对于params最大化目标函数（默认值：False），而不是最小化。

• capturable (bool, 可选) – 是否可以将此实例安全地捕获到 CUDA 图中。设置为 True 可能会影响未进行图绘制的性能，因此如果你不打算对此实例进行图绘制，请保持为 False（默认值：False）

• 可微分 (bool, 可选) – 在训练过程中，autograd 是否应通过优化器步骤进行。如果设置为 False，则 step() 函数将在 torch.no_grad() 上下文中运行。将此参数设为 True 可能会影响性能，因此如果你不打算在此实例中使用 autograd，请保持默认值 False。

• fused (bool, 可选) – 是否使用融合实现。当前支持的数据类型包括torch.float64、torch.float32、torch.float16和torch.bfloat16。（默认值：None）

注意

foreach 和融合实现通常比单个循环和单张量实现更快，其中融合在理论上通过垂直和水平融合实现了最快的速度。因此，如果用户没有指定任何标志（即，当 foreach = fused = None 时），我们将在所有张量都在 CUDA 上的情况下默认使用 foreach 实现。为什么不直接使用融合？因为融合实现相对较新，我们需要给它足够的磨合时间。要启用融合，请将 fused 设为 True。要强制运行循环实现，请将 foreach 或 fused 设为 False。

注意

MPS中的Adam和AdamW的原型实现支持torch.float32和torch.float16。

add_param_group(param_group)

将一个参数组添加到Optimizer的param_groups中。

在微调预训练网络时，这非常有用，因为可以将冻结的层设为可训练状态，并随着训练进度将其添加到Optimizer中。

参数

param_group (dict) – 指定需要优化的张量及其特定的优化选项。

load_state_dict(state_dict)

加载优化器的状态。

参数

state_dict (dict) – 优化器的状态。应为调用state_dict() 返回的字典对象。

register_load_state_dict_post_hook(hook, prepend=False)

注册一个在调用load_state_dict()之后触发的 post-hook。它应具有以下签名：

hook(optimizer) -> None

optimizer 参数是指正在使用的优化器实例。

在调用 self.load_state_dict() 之后，该钩子将被调用，并使用参数 self。注册的钩子可以在 load_state_dict 加载了 state_dict 之后执行后处理。

参数

• hook (Callable) – 用户自定义的要注册的钩子。

• prepend (bool) – 如果为 True，则提供的后置钩子 hook 将在调用 load_state_dict 时，在所有已注册的后置钩子之前触发。否则，该钩子将在所有已注册的后置钩子之后触发。（默认值：False）

返回值

一个可以通过调用 handle.remove() 移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemoveableHandle

register_load_state_dict_pre_hook(hook, prepend=False)

注册一个在调用load_state_dict()之前被调用的load_state_dict预处理钩子。该钩子应具有以下签名：

hook(optimizer, state_dict) -> state_dict or None

optimizer 参数是正在使用的优化器实例，而 state_dict 参数是用户传递给 load_state_dict 的 state_dict 浅拷贝。钩子可以就地修改 state_dict 或者可选地返回一个新的 state_dict。如果返回了一个新的 state_dict，它将被加载到优化器中。

该钩子在调用 self.load_state_dict() 之前会被调用，并使用参数 self 和 state_dict。注册的钩子可以在 load_state_dict 调用前用于执行预处理。

参数

• hook (Callable) – 用户自定义的要注册的钩子。

• prepend (bool) – 如果为 True，提供的 hook 将在调用 load_state_dict 之前触发所有已注册的预处理钩子。否则，该 hook 将在所有已注册的预处理钩子之后触发。（默认值：False）

返回值

一个可以通过调用 handle.remove() 移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemoveableHandle

register_state_dict_post_hook(hook, prepend=False)

注册一个状态字典的后置钩子，在调用 state_dict() 之后会被调用。

它应具有如下签名：

hook(optimizer, state_dict) -> state_dict or None

在生成self的state_dict之后，该钩子将被调用，并使用self和state_dict作为参数。钩子可以就地修改state_dict或选择返回一个新的state_dict。注册的钩子可以在返回之前对state_dict执行后处理。

参数

• hook (Callable) – 用户自定义的要注册的钩子。

• prepend (bool) – 如果为 True，则提供的 post hook 将在所有已注册的 post-hooks 之前触发。否则，提供的 hook 将在所有已注册的 post-hooks 之后触发。（默认值：False）

返回值

一个可以通过调用 handle.remove() 移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemoveableHandle

register_state_dict_pre_hook(hook, prepend=False)

注册一个在调用 state_dict() 之前执行的状态字典预处理钩子。

它应具有如下签名：

hook(optimizer) -> None

optimizer 参数是正在使用的优化器实例。在调用 self.state_dict() 方法之前，会将 self 作为参数传递给钩子。注册的钩子可以在调用 state_dict 之前执行预处理。

参数

• hook (Callable) – 用户自定义的要注册的钩子。

• prepend (bool) – 如果为 True，则提供的 hook 将在所有已注册的预处理钩子之前触发。否则，它将在所有已注册的预处理钩子之后触发。（默认值：False）

返回值

一个可以通过调用 handle.remove() 移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemoveableHandle

register_step_post_hook(hook)

注册一个在优化器步骤之后调用的后处理钩子。

它应具有如下签名：

hook(optimizer, args, kwargs) -> None

optimizer 参数是指正在使用的优化器实例。

参数

hook (Callable) – 用户自定义的要注册的钩子。

返回值

一个可以通过调用 handle.remove() 移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

register_step_pre_hook(hook)

注册一个在优化器步骤执行前被调用的预处理钩子。

它应具有如下签名：

hook(optimizer, args, kwargs) -> None or modified args and kwargs

optimizer 参数是指正在使用的优化器实例。如果预处理钩子修改了 args 和 kwargs，则会返回一个新的包含修改后参数的元组，即新的 new_args 和 new_kwargs。

参数

hook (Callable) – 用户自定义的要注册的钩子。

返回值

一个可以通过调用 handle.remove() 移除添加的钩子的句柄

返回类型

torch.utils.hooks.RemovableHandle

state_dict()

以dict形式返回优化器的状态。

它包含了两个条目：

• state: 一个包含当前优化状态的 Dict。

  虽然在不同的优化器类之间存在差异，但也有一些共同的特点。例如，状态是按参数进行保存的，并且不会保存参数本身。state 是一个字典，将参数ID映射到包含每个参数对应状态信息的字典。

• param_groups: 包含所有参数组的列表，每个参数组内

  参数组是一个字典。每个参数组包含了特定于优化器的元数据（如学习率和权重衰减），以及该组内参数的ID列表。

注意：参数 ID 虽然看起来像索引，但实际上只是用于将状态与 param_group 关联的标识符。在从 state_dict 加载时，优化器会按顺序将 param_group 中的 params（整数 ID）和优化器中的 param_groups （实际的 nn.Parameters）进行 zip 操作以匹配状态，而无需额外验证。

返回的 state 字典可能看起来像这样：

{
    'state': {
        0: {'momentum_buffer': tensor(...), ...},
        1: {'momentum_buffer': tensor(...), ...},
        2: {'momentum_buffer': tensor(...), ...},
        3: {'momentum_buffer': tensor(...), ...}
    },
    'param_groups': [
        {
            'lr': 0.01,
            'weight_decay': 0,
            ...
            'params': [0]
        },
        {
            'lr': 0.001,
            'weight_decay': 0.5,
            ...
            'params': [1, 2, 3]
        }
    ]
}

返回类型

Dict[str, Any]

step(closure=None)[源代码]

执行一步优化。

参数

closure (Callable, 可选) – 一个可选的闭包，用于重新评估模型并返回损失值。

zero_grad(set_to_none=True)

将所有优化的 torch.Tensor 的梯度重置。

参数

set_to_none (bool) – 代替将梯度设置为零，将其设置为 None。这通常会减少内存占用，并可以适度提高性能。然而，它改变了某些行为：1. 当用户尝试访问梯度并执行手动操作时，None 属性或全为0的 Tensor 的表现不同。2. 如果用户请求 zero_grad(set_to_none=True) 后跟一次反向传播，对于未接收到梯度的参数，.grads 保证为 None。3. torch.optim 优化器在梯度为0或 None 的情况下有不同的行为：在一个情况下它使用零梯度进行步骤，在另一个情况下则完全跳过该步骤。