SparseAdam

类torch.optim.SparseAdam(params, lr=0.001, betas=(0.9, 0.999), eps=1e-08, maximize=False)[源代码]

SparseAdam 实现了适用于稀疏梯度的 Adam 算法的一个版本。由于实现限制，目前 SparseAdam 仅适用于特定的一小部分用例：参数为稠密布局而梯度为稀疏布局的情况。这种情况发生在模块反向传播时生成已经处于稀疏布局的梯度。一个以这种方式工作的示例神经网络模块是 nn.Embedding(sparse=True)。

SparseAdam 通过屏蔽与梯度中零值相对应的参数和动量更新来近似 Adam 算法。而 Adam 算法则会根据梯度的所有值更新第一个动量、第二个动量以及参数，SparseAdam 则只对非零梯度值对应的动量和参数进行更新。

一种简单的方式来理解预期实现是：

创建一个稀疏梯度中非零值的掩码。例如，如果梯度为 [0, 5, 0, 0, 9]，那么掩码将会是 [0, 1, 0, 0, 1]。
应用此掩码来处理运行时刻的数据，并仅对非零值进行计算。
将此掩码应用于参数，并仅对非零值进行更新。

实际上，我们通过使用稀疏布局张量来优化这个近似值，这意味着未实际化的梯度越多，优化性能越好。由于我们依赖于稀疏布局张量，因此我们可以推断出在稀疏布局中任何已实现的值都是非零的，并且不会验证所有值是否为零。重要的是不要将语义上稀疏的张量（其中许多值为零）与稀疏布局张量（.is_sparse 返回 True 的张量）混淆。SparseAdam 近似值是针对语义上的稀疏张量设计的，而稀疏布局只是一个实现细节。更清晰的实现方式是使用 MaskedTensors，但这些功能还在实验阶段。

注意

如果你认为你的梯度在语义上是稀疏的（但没有采用稀疏布局），这个变体可能不适合你。理想情况下，你应该避免生成任何疑似稀疏的数据，因为将所有密集布局的梯度转换为稀疏布局可能会抵消性能提升带来的好处。在这种情况下，使用Adam可能是最好的替代方案，除非你可以轻松地让你的模块输出类似于nn.Embedding(sparse=True) 的稀疏梯度。如果你坚持要转换你的梯度，你可以在调用.step() 之前手动覆盖参数的.grad 字段以使用其稀疏等价物。

参数

params (可迭代对象) – 需要优化的参数或定义参数组的字典的可迭代对象
lr (float 或 Tensor, 可选) – 学习率（默认值：1e-3）
betas (Tuple[float, float], optional) – 用于计算梯度及其平方的滑动平均系数（默认值：(0.9, 0.999)）
eps (float, 可选) – 用于提高数值稳定性的添加到分母中的项（默认值：1e-8）
maximize (bool, optional) – 是否相对于params最大化目标函数（默认值：False），而不是最小化。

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 的情况下有不同的行为：在一个情况下它使用零梯度进行步骤，在另一个情况下则完全跳过该步骤。