SparseAdam¶
- class torch.optim.SparseAdam(params, lr=0.001, betas=(0.9, 0.999), eps=1e-08, maximize=False)[source][source]¶
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(浮点数,张量,可选)- 学习率(默认:1e-3)
betas(Tuple[float, float],可选)- 用于计算梯度及其平方的运行平均值的系数(默认:(0.9, 0.999))
eps(浮点数,可选)- 添加到分母的项,以提高数值稳定性(默认:1e-8)
maximize (bool, 可选) – 相对于 params 最大化目标,而不是最小化(默认:False)
- add_param_group(param_group)[source]
向
Optimizer
s 的 param_groups 中添加一个参数组。这在微调预训练网络时可能很有用,因为冻结的层可以在训练过程中变为可训练并添加到
Optimizer
中。- 参数:
param_group (dict) - 指定哪些张量应与特定优化选项一起优化。
- load_state_dict(state_dict)[source]¶
加载优化器状态。
- 参数:
state_dict (dict) – 优化器状态。应该是从调用
state_dict()
返回的对象。
注意
参数的名称(如果它们存在于每个参数组在
state_dict()
的“param_names”键下)将不会影响加载过程。如果要使用参数的名称进行自定义情况(例如,当加载的状态字典中的参数与初始化优化器中的参数不同时),则需要实现自定义的register_load_state_dict_pre_hook
以相应地适应加载的字典。如果加载的状态字典中存在param_names
,则将保存并覆盖优化器状态中当前存在的名称。如果它们不在加载的状态字典中,则优化器的param_names
将保持不变。
- register_load_state_dict_post_hook(hook, prepend=False)[source]
注册一个在调用
load_state_dict()
之后调用的 load_state_dict 后钩子。它应该具有以下签名:hook(optimizer) -> None
第 0#个参数是正在使用的优化器实例。
在对
self
调用load_state_dict
之后,将使用参数self
调用钩子。已注册的钩子可以在load_state_dict
加载state_dict
之后执行后处理。- 参数:
hook(可调用对象)- 用户定义的将被注册的钩子。
prepend (bool) – 如果为 True,提供的后处理
hook
将在所有已注册的后处理钩子load_state_dict
之前触发。否则,提供的hook
将在所有已注册的后处理钩子之后触发。(默认:False)
- 返回值:
一个可以用来通过调用
handle.remove()
移除已添加钩子的句柄。- 返回类型:
torch.utils.hooks.RemoveableHandle
- register_load_state_dict_pre_hook(hook, prepend=False)[source]¶
注册一个在调用
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,它将被用于加载到优化器中。在调用
self
上的load_state_dict
之前,将使用参数self
和state_dict
调用钩子。已注册的钩子可用于在执行load_state_dict
调用之前执行预处理。- 参数:
hook(可调用对象)- 用户定义的将被注册的钩子。
prepend (bool) – 如果为 True,提供的预
hook
将在所有已注册的预钩子load_state_dict
之前触发。否则,提供的hook
将在所有已注册的预钩子之后触发。(默认:False)
- 返回值:
一个可以用来通过调用
handle.remove()
移除已添加钩子的句柄。- 返回类型:
torch.utils.hooks.RemoveableHandle
- register_state_dict_post_hook(hook, prepend=False)[source]¶
注册一个状态字典后钩子,该钩子将在调用
state_dict()
之后被调用。它应该具有以下签名:
hook(optimizer, state_dict) -> state_dict or None
在生成
self
上的state_dict
之后,将以参数self
和state_dict
调用钩子。钩子可以就地修改状态字典或可选地返回一个新的状态字典。已注册的钩子可用于在返回之前对state_dict
执行后处理。- 参数:
hook(可调用对象)- 用户定义的将被注册的钩子。
prepend (bool) – 如果为 True,提供的后处理
hook
将在所有已注册的后处理钩子state_dict
之前触发。否则,提供的hook
将在所有已注册的后处理钩子之后触发。(默认:False)
- 返回值:
一个可以用来通过调用
handle.remove()
移除已添加钩子的句柄。- 返回类型:
torch.utils.hooks.RemoveableHandle
- register_state_dict_pre_hook(hook, prepend=False)[source]¶
注册一个在调用
state_dict()
之前被调用的状态字典预钩子。它应该具有以下签名:
hook(optimizer) -> None
The
optimizer
argument is the optimizer instance being used. The hook will be called with argumentself
before callingstate_dict
onself
. The registered hook can be used to perform pre-processing before thestate_dict
call is made.- 参数:
hook(可调用对象)- 用户定义的将被注册的钩子。
prepend (bool) – 如果为 True,提供的预
hook
将在所有已注册的预钩子state_dict
之前触发。否则,提供的hook
将在所有已注册的预钩子之后触发。(默认:False)
- 返回值:
一个可以用来通过调用
handle.remove()
移除已添加钩子的句柄。- 返回类型:
torch.utils.hooks.RemoveableHandle
- register_step_post_hook(hook)[source]¶
注册优化器步骤后钩子,该钩子在优化器步骤之后将被调用。
它应该具有以下签名:
hook(optimizer, args, kwargs) -> None
第 0#个参数是正在使用的优化器实例。
- 参数:
hook(可调用对象)- 用户定义的将被注册的钩子。
- 返回值:
一个可以用来通过调用
handle.remove()
移除已添加钩子的句柄。- 返回类型:
torch.utils.hooks.RemovableHandle
- register_step_pre_hook(hook)[source]¶
注册优化器步骤预钩子,该钩子将在优化器步骤之前被调用。
它应该具有以下签名:
hook(optimizer, args, kwargs) -> None or modified args and kwargs
第 0#参数是正在使用的优化器实例。如果 args 和 kwargs 被预钩子修改,则返回包含新_args 和新_kwargs 的新值元组。
- 参数:
hook(可调用对象)- 用户定义的将被注册的钩子。
- 返回值:
一个可以用来通过调用
handle.remove()
移除已添加钩子的句柄。- 返回类型:
torch.utils.hooks.RemovableHandle
- state_dict()[source]¶
返回优化器的状态作为
dict
。它包含两个条目:
-
state
:一个字典,包含当前的优化状态。其内容
优化器类之间有所不同,但有一些共同特征。例如,状态是按参数保存的,而参数本身不保存。state
是一个字典,将参数 ID 映射到与每个参数对应的状态的字典。
-
- 包含所有参数组的列表
参数组是一个字典。每个参数组包含特定于优化器的元数据,例如学习率和权重衰减,以及组中参数的 ID 列表。如果参数组使用named_parameters()
初始化,则名称内容也将保存在状态字典中。
注意:参数 ID 看起来像索引,但它们只是将状态与 param_group 关联的 ID。在从 state_dict 加载时,优化器将 param_group
params
(整数 ID)和优化器param_groups
(实际的nn.Parameter
s)进行 zip,以便匹配状态,无需额外验证。返回的状态字典可能看起来像这样:
{ '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]
- zero_grad(set_to_none=True)[source]
重置所有优化
torch.Tensor
的梯度。- 参数:
set_to_none (布尔值) – 在设置为零而不是设置为零的情况下,将梯度设置为 None。这通常会有更低的内存占用,并且可以适度提高性能。但是,它会改变某些行为。例如:1. 当用户尝试访问梯度并对其进行手动操作时,None 属性或充满 0s 的张量将会有不同的行为。2. 如果用户请求
zero_grad(set_to_none=True)
然后进行反向传播,.grad
确保对于未接收梯度的参数将是 None。3.torch.optim
优化器在梯度为 0 或 None 时具有不同的行为(在一种情况下它使用梯度为 0 执行步骤,在另一种情况下它完全跳过步骤)。