Adafactor

classtorch.optim.Adafactor(params, lr=0.01, beta2_decay=-0.8, eps=(None, 0.001), d=1.0, weight_decay=0.0, *, foreach=None, maximize=False)

实现Adafactor算法。

$\begin{aligned} &\rule{110mm}{0.4pt} \\ &\textbf{input} : \gamma \text{(lr)}, \: \tau \text{(}\beta_2\text{ decay)}, \: \theta_0 \text{(params)}, \: f(\theta) \text{(objective)}, \\ &\hspace{15mm} \: \epsilon_1, \epsilon_2 \text{ (epsilons)}, \: d \text{(clipping threshold)}, \\ &\hspace{15mm} \: \lambda \text{(weight decay)}, \: \textit{maximize} \\ &\textbf{initialize} : \: R_0 \leftarrow 0 \text{ (second moment row factor)}, \\ &\hspace{23mm} \: C_0 \leftarrow 0 \text{ (second moment col factor)}, \\ &\hspace{23mm} \: \widehat{V}_0 \leftarrow 0 \text{ (second moment for vectors)} \\[-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}\widehat{\beta}_{2_t} \leftarrow 1 - t^{\tau} \\ &\hspace{5mm}\rho_t \leftarrow min(lr, \frac{1}{\sqrt{t}}) \\ &\hspace{5mm}\alpha_t \leftarrow max(\epsilon_2, \text{RMS}(\theta_{t-1}))\rho_t \\ &\hspace{5mm}\theta_t \leftarrow \theta_{t-1} - \gamma \lambda \theta_{t-1} \\ &\hspace{5mm}\textbf{if} \: \text{dim}(G_t) > 1: \\ &\hspace{10mm}R_t \leftarrow \widehat{\beta}_{2_t}R_{t-1}+ (1-\widehat{\beta}_{2_t})(G_t \odot G_t) \cdot 1_m \\ &\hspace{10mm}C_t \leftarrow \widehat{\beta}_{2_t}C_{t-1}+ (1-\widehat{\beta}_{2_t}) 1^\top_n \cdot (G_t \odot G_t) \\ &\hspace{10mm}\widehat{V}_t \leftarrow \frac{R_t \cdot C_t}{max(1^\top_n \cdot R_t, \epsilon_1)} \\ &\hspace{5mm}\textbf{else} \\ &\hspace{10mm}\widehat{V}_t \leftarrow \widehat{\beta}_{2_t}\widehat{V}_{t-1}+ (1-\widehat{\beta}_{2_t}) \cdot (G_t \odot G_t) \\ &\hspace{5mm}U_t \leftarrow \frac{G_t}{max(\sqrt{\widehat{V}_t}, \epsilon_1)} \\ &\hspace{5mm}\widehat{U}_t \leftarrow \frac{U_t}{max(1, \frac{\text{RMS}(U_t)}{d})} \\ &\hspace{5mm}\theta_t \leftarrow \theta_{t-1} - \alpha_t \widehat{U}_t \\ &\rule{110mm}{0.4pt} \\[-1.ex] &\bf{return} \: \theta_t \\[-1.ex] &\rule{110mm}{0.4pt} \\[-1.ex] \end{aligned}$

关于算法的更多细节，请参阅 Adafactor: 自适应学习率，亚线性内存成本。

参数

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

• lr (float, Tensor, 可选) – 与其它优化器不同，Adafactor 不需要学习率，并且 Shazeer、Noam 和 Mitchell Stern 完全不使用 lr。本实现中使用 lr 来应用权重衰减以及作为相对步长 rho_t 的最大值，这不同于论文中的做法。需要注意的是，在论文中，常数 0.01 被用作相对步长的最大值，默认情况下我们将其设置为 0.01。(默认值: 1e-2)

• beta2_decay (float, 可选) – beta2 的衰减率。通常，beta2 是用于计算梯度平方滑动平均的系数。(默认值: -0.8)

• eps (Tuple[float, float], 可选) – epslion1 是添加到更新计算分母中的项，以提高数值稳定性。这与论文中描述的算法有所不同！有关更多详细信息，请参见下方注释。epslion2 用于在应用参数缩放时避免权重更新过小。（默认值：(None, 1e-3))

• d (float, 可选) – 剪辑阈值，用于防止更新过大。

• weight_decay (float, optional) – 权重衰减系数，默认值为 1e-2

• foreach (bool, 可选) – 是否使用优化器的 foreach 实现。注意，由于中间结果是张量列表而不是单个张量，foreach 实现比 for 循环版本多占用约 sizeof(params) 的峰值内存。通常在内存受限的情况下会使用 Adafactor，默认情况下它会采用较慢的单个张量 for 循环实现，除非明确设置此标志为 True。与其他优化器不同，它们默认会在 CUDA 上尝试使用 foreach 实现以加快运行时间。（默认值：None）

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

注意

Adafactor 的实现与其他实现（如 Shazeer、Noam 和 Mitchell Stern 以及某些其他框架）在学习率和 $\epsilon_1$ 的使用上有细微差别。

关于学习率超参数：Shazeer、Noam 和 Mitchell Stern 不使用 lr，因为他们使用的算法通过$\rho_t$和更新裁剪来控制步长。

此实现允许lr影响$\rho_t$的最大值:

$\begin{aligned} &\hspace{5mm}\rho_t \leftarrow min(lr, \frac{1}{\sqrt{t}}) \end{aligned}$

这与Shazeer、Noam和Mitchell Stern的不同之处在于，他们将0.01设置为$\rho_t$的最大值。

$\begin{aligned} &\hspace{5mm}\rho_t \leftarrow min(0.01, \frac{1}{\sqrt{t}}) \end{aligned}$

Shazeer 和 Mitchell Stern 没有对权重衰减的计算方式提出具体意见，因此我们使用学习率作为解耦权重衰减的系数，类似于《解耦权重衰减正则化》中的建议。

关于使用$\epsilon_1$：实现尝试复制Shazeer、Noam和Mitchell Stern的意图，即在平方梯度变小时将$\epsilon_1$用作稳定项。

这个稳定性可以表示为

$\begin{aligned} &\hspace{5mm}R_t \leftarrow \widehat{\beta}_{2_t}R_{t-1}+ (1-\widehat{\beta}_{2_t})(G_t \odot G_t + 1_n \cdot 1^\top_m) \cdot 1_m \\ &\hspace{5mm}C_t \leftarrow \widehat{\beta}_{2_t}C_{t-1}+ (1-\widehat{\beta}_{2_t}) 1^\top_n \cdot (G_t \odot G_t + 1_n \cdot 1^\top_m) \\ &\hspace{5mm}\widehat{V}_t \leftarrow \frac{R_t \cdot C_t}{max(1^\top_n \cdot R_t, \epsilon_1)} \\ &\hspace{5mm}U_t \leftarrow \frac{G_t}{max(\sqrt{\widehat{V}_t}, \epsilon_1)} \\ \end{aligned}$

其中，梯度平方的行因子 $R_t$ 和列因子 $C_t$ 保持不变，在方差估计 $\widehat{V}_t$ 的最终计算以及更新 $U_t$ 中应用 $\epsilon_1$。

这与 Shazeer、Noam 和 Mitchell Stern 以及其他的框架不同，这些框架将 $\epsilon_1$ 应用于平方梯度的行和列因子，但在后续计算中不应用：

$\begin{aligned} &\hspace{5mm}R_t \leftarrow \widehat{\beta}_{2_t}R_{t-1}+ (1-\widehat{\beta}_{2_t})(G_t \odot G_t + \epsilon_1 1_n \cdot 1^\top_m) \cdot 1_m \\ &\hspace{5mm}C_t \leftarrow \widehat{\beta}_{2_t}C_{t-1}+ (1-\widehat{\beta}_{2_t}) 1^\top_n \cdot (G_t \odot G_t + \epsilon_1 1_n \cdot 1^\top_m) \\ &\hspace{5mm}\widehat{V}_t \leftarrow \frac{R_t \cdot C_t}{1^\top_n \cdot R_t} \\ &\hspace{5mm}U_t \leftarrow \frac{G_t}{\sqrt{\widehat{V}_t}} \\ \end{aligned}$

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