SparseAdam

class torch.optim.SparseAdam(params, lr=0.001, betas=(0.9, 0.999), eps=1e-08, maximize=False)

SparseAdam implements a masked version of the Adam algorithm suitable for sparse gradients. Currently, due to implementation constraints (explained below), SparseAdam is only intended for a narrow subset of use cases, specifically parameters of a dense layout with gradients of a sparse layout. This occurs in a special case where the module backwards produces grads already in a sparse layout. One example NN module that behaves as such is nn.Embedding(sparse=True).

SparseAdam approximates the Adam algorithm by masking out the parameter and moment updates corresponding to the zero values in the gradients. Whereas the Adam algorithm will update the first moment, the second moment, and the parameters based on all values of the gradients, SparseAdam only updates the moments and parameters corresponding to the non-zero values of the gradients.

A simplified way of thinking about the intended implementation is as such

1. Create a mask of the non-zero values in the sparse gradients. For example, if your gradient looks like [0, 5, 0, 0, 9], the mask would be [0, 1, 0, 0, 1].

2. Apply this mask over the running moments and do computation on only the non-zero values.

3. Apply this mask over the parameters and only apply an update on non-zero values.

In actuality, we use sparse layout Tensors to optimize this approximation, which means the more gradients that are masked by not being materialized, the more performant the optimization. Since we rely on using sparse layout tensors, we infer that any materialized value in the sparse layout is non-zero and we do NOT actually verify that all values are not zero! It is important to not conflate a semantically sparse tensor (a tensor where many of its values are zeros) with a sparse layout tensor (a tensor where .is_sparse returns True). The SparseAdam approximation is intended for semantically sparse tensors and the sparse layout is only a implementation detail. A clearer implementation would be to use MaskedTensors, but those are experimental.

注意

If you suspect your gradients are semantically sparse (but do not have sparse layout), this variant may not be the best for you. Ideally, you want to avoid materializing anything that is suspected to be sparse in the first place, since needing to convert all your grads from dense layout to sparse layout may outweigh the performance gain. Here, using Adam may be the best alternative, unless you can easily rig up your module to output sparse grads similar to nn.Embedding(sparse=True). If you insist on converting your grads, you can do so by manually overriding your parameters’ .grad fields with their sparse equivalents before calling .step().

参数

• params (iterable) – 要优化的参数或命名参数的迭代器，或者是定义参数组的字典的迭代器。使用命名参数时，所有组中的所有参数都应该命名。

• lr (float, Tensor, optional) – 学习率 (默认: 1e-3)

• betas (Tuple[float, float], optional) – 用于计算梯度及其平方的运行平均值的系数 (默认: (0.9, 0.999))

• eps (float, optional) – 添加到分母中的项，以提高数值稳定性 (默认: 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() 返回的对象。

警告

请确保在初始化 torch.optim.lr_scheduler.LRScheduler 后调用此方法，因为提前调用它会覆盖加载的学习率。

注意

参数的名称（如果它们存在于 state_dict() 中每个参数组的“param_names”键下）不会影响加载过程。对于自定义用例（例如，当加载的状态字典中的参数与优化器中初始化的参数不同时），可以使用自定义的 register_load_state_dict_pre_hook 来适应加载的字典。如果加载的状态字典的 param_groups 中存在 param_names，它们将被保存并覆盖优化器状态中当前的名称。如果它们不存在于加载的状态字典中，优化器的 param_names 将保持不变。

示例

>>> model = torch.nn.Linear(10, 10)
>>> optim = torch.optim.SGD(model.parameters(), lr=3e-4)
>>> scheduler1 = torch.optim.lr_scheduler.LinearLR(
...     optim,
...     start_factor=0.1,
...     end_factor=1,
...     total_iters=20,
... )
>>> scheduler2 = torch.optim.lr_scheduler.CosineAnnealingLR(
...     optim,
...     T_max=80,
...     eta_min=3e-5,
... )
>>> lr = torch.optim.lr_scheduler.SequentialLR(
...     optim,
...     schedulers=[scheduler1, scheduler2],
...     milestones=[20],
... )
>>> lr.load_state_dict(torch.load("./save_seq.pt"))
>>> # now load the optimizer checkpoint after loading the LRScheduler
>>> optim.load_state_dict(torch.load("./save_optim.pt"))

register_load_state_dict_post_hook(hook, prepend=False)

注册一个 load_state_dict 后置钩子，它将在调用 load_state_dict() 之后调用。它应该具有以下签名

hook(optimizer) -> None

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

调用 load_state_dict 后，将使用参数 self 调用该钩子。已注册的钩子可用于在 load_state_dict 加载 state_dict 后执行后处理。

参数

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

• prepend (bool) – 如果为 True，则提供的 load_state_dict 后置 hook 将在所有已注册的 load_state_dict 后置钩子之前执行。否则，提供的 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

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

在调用 load_state_dict 之前，将使用参数 self 和 state_dict 调用该钩子。已注册的钩子可用于在调用 load_state_dict 之前执行预处理。

参数

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

• prepend (bool) – 如果为 True，则提供的 load_state_dict 前置 hook 将在所有已注册的 load_state_dict 前置钩子之前执行。否则，提供的 hook 将在所有已注册的 load_state_dict 前置钩子之后执行。(默认: False)

返回

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

返回类型

torch.utils.hooks.RemoveableHandle

register_state_dict_post_hook(hook, prepend=False)

注册一个 state_dict 后置钩子，它将在调用 state_dict() 之后调用。

它应具有以下签名

hook(optimizer, state_dict) -> state_dict or None

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

参数

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

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

返回

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

返回类型

torch.utils.hooks.RemoveableHandle

register_state_dict_pre_hook(hook, prepend=False)

注册一个 state dict pre-hook，在调用 state_dict() 之前调用。

它应具有以下签名

hook(optimizer) -> None

optimizer 参数是正在使用的优化器实例。hook 将在调用 `self` 上的 `state_dict` 之前，并使用 `self` 作为参数进行调用。注册的 hook 可用于在调用 `state_dict` 之前执行预处理。

参数

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

• prepend (bool) – 如果为 True，则提供的 pre hook 将在 state_dict 上所有已注册的 pre-hook 之前触发。否则，提供的 hook 将在 state_dict 上所有已注册的 pre-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 被 pre-hook 修改，则转换后的值将作为包含 new_args 和 new_kwargs 的元组返回。

参数

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

返回

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

返回类型

torch.utils.hooks.RemovableHandle

state_dict()

将优化器的状态作为 dict 返回。

它包含两个条目

• state：一个包含当前优化状态的 Dict。其内容

  优化器类之间会有所不同，但一些共同的特性是适用的。例如，状态是按参数保存的，而参数本身不保存。`state` 是一个字典，将参数 ID 映射到包含每个参数对应状态的字典。

• param_groups：一个包含所有参数组的 List，其中每个

  参数组是一个字典。每个参数组包含优化器特有的元数据，如学习率和权重衰减，以及组内参数的 ID 列表。如果一个参数组是用 `named_parameters()` 初始化，那么名称内容也会保存在 state_dict 中。

注意：参数 ID 可能看起来像索引，但它们只是将状态与参数组关联起来的 ID。在从 state_dict 加载时，优化器将参数组的 `params`（整数 ID）和优化器的 `param_groups`（实际的 `nn.Parameter`）按顺序组合起来，以匹配状态，而无需额外的验证。

返回的状态字典可能看起来像

{
    '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]
            'param_names' ['param0']  (optional)
        },
        {
            'lr': 0.001,
            'weight_decay': 0.5,
            ...
            'params': [1, 2, 3]
            'param_names': ['param1', 'layer.weight', 'layer.bias'] (optional)
        }
    ]
}

返回类型

dict[str, Any]

step(closure=None)

执行单个优化步骤。

参数

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

zero_grad(set_to_none=True)

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

参数

set_to_none (bool) – 与其设置为零，不如将 grad 设置为 None。这通常会降低内存占用，并可能适度提高性能。但是，它会改变某些行为。例如：1. 当用户尝试访问 grad 并对其进行手动操作时，None 属性或全零的 Tensor 的行为将不同。2. 如果用户请求 zero_grad(set_to_none=True)，然后进行反向传播，对于未接收到梯度的参数，`.grad` 保证为 None。3. `torch.optim` 优化器在 grad 为 0 或 None 时具有不同的行为（一种情况是使用 grad 0 执行步长，另一种情况是完全跳过步长）。