mindspore.nn.AdamOffload

class mindspore.nn.AdamOffload(params, learning_rate=0.001, beta1=0.9, beta2=0.999, eps=1e-08, use_locking=False, use_nesterov=False, weight_decay=0.0, loss_scale=1.0)[源代码]

此优化器在主机CPU上运行Adam优化算法,设备上仅执行网络参数的更新,最大限度地降低内存成本。虽然会增加性能开销,但优化器可被用于运行更大的模型。

Adam算法参见 Adam: A Method for Stochastic Optimization

更新公式如下:

\[\begin{split}\begin{array}{ll} \\ m_{t+1} = \beta_1 * m_{t} + (1 - \beta_1) * g \\ v_{t+1} = \beta_2 * v_{t} + (1 - \beta_2) * g * g \\ l = \alpha * \frac{\sqrt{1-\beta_2^t}}{1-\beta_1^t} \\ w_{t+1} = w_{t} - l * \frac{m_{t+1}}{\sqrt{v_{t+1}} + \epsilon} \end{array}\end{split}\]

\(m\) 代表第一个矩向量 moment1\(v\) 代表第二个矩向量 moment2\(g\) 代表 gradients\(l\) 代表缩放因子,\(\beta_1,\beta_2\) 代表 beta1beta2\(t\) 代表当前step,\(beta_1^t\)\(beta_2^t\) 代表 beta1_powerbeta2_power\(\alpha\) 代表 learning_rate\(w\) 代表 params\(\epsilon\) 代表 eps

Note

此优化器目前仅支持图模式。

在参数未分组时,优化器配置的 weight_decay 应用于名称不含”beta”或”gamma”的网络参数。

参数分组情况下,可以分组调整权重衰减策略。

分组时,每组网络参数均可配置 weight_decay ,若未配置,则该组网络参数使用优化器中配置的 weight_decay

参数:

  • params (Union[list[Parameter], list[dict]]) - 必须是 Parameter 组成的列表或字典组成的列表。当列表元素是字典时,字典的键可以是”params”、”lr”、”weight_decay”、和”order_params”:

    • params - 必填。当前组别的权重,该值必须是 Parameter 列表。

    • lr - 可选。如果键中存在”lr”,则使用对应的值作为学习率。如果没有,则使用优化器中的参数 learning_rate 作为学习率。支持固定和动态学习率。

    • weight_decay - 可选。如果键中存在”weight_decay”,则使用对应的值作为权重衰减值。如果没有,则使用优化器中配置的 weight_decay 作为权重衰减值。 值得注意的是, weight_decay 可以是常量,也可以是Cell类型。Cell类型的weight decay用于实现动态weight decay算法。这和动态学习率相似。 用户需要自定义一个输入为global step的weight_decay_schedule。在训练的过程中,优化器会调用WeightDecaySchedule的实例来获取当前step的weight decay值。

    • order_params - 可选。值的顺序是参数更新的顺序。当使用参数分组功能时,通常使用该配置项保持 parameters 的顺序以提升性能。如果键中存在”order_params”,则会忽略该组配置中的其他键。”order_params”中的参数必须在某一组 params 参数中。

  • learning_rate (Union[float, Tensor, Iterable, LearningRateSchedule]): 默认值:1e-3。

    • float - 固定的学习率。必须大于等于零。

    • int - 固定的学习率。必须大于等于零。整数类型会被转换为浮点数。

    • Tensor - 可以是标量或一维向量。标量是固定的学习率。一维向量是动态的学习率,第i步将取向量中第i个值作为学习率。

    • Iterable - 动态的学习率。第i步将取迭代器第i个值作为学习率。

    • LearningRateSchedule - 动态的学习率。在训练过程中,优化器将使用步数(step)作为输入,调用 LearningRateSchedule 实例来计算当前学习率。

  • beta1 (float) - moment1 的指数衰减率。参数范围(0.0,1.0)。默认值:0.9。

  • beta2 (float) - moment2 的指数衰减率。参数范围(0.0,1.0)。默认值:0.999。

  • eps (float) - 将添加到分母中,以提高数值稳定性。必须大于0。默认值:1e-8。

  • use_locking (bool) - 是否对参数更新加锁保护。如果为True,则 wmv 的更新将受到锁保护。如果为False,则结果不可预测。默认值:False。

  • use_nesterov (bool) - 是否使用Nesterov Accelerated Gradient (NAG)算法更新梯度。如果为True,使用NAG更新梯度。如果为False,则在不使用NAG的情况下更新梯度。默认值:False。

  • weight_decay (Union[float, int, Cell]) - 权重衰减(L2 penalty)。默认值:0.0。

    • float: 固定值,必须大于或者等于0。

    • int: 固定值,必须大于或者等于0,会被转换成float。

    • Cell: 动态weight decay。在训练过程中,优化器会使用步数(step)作为输入,调用该Cell实例来计算当前weight decay值。

  • loss_scale (float) - 梯度缩放系数,必须大于0。如果 loss_scale 是整数,它将被转换为浮点数。通常使用默认值,仅当训练时使用了 FixedLossScaleManager,且 FixedLossScaleManagerdrop_overflow_update 属性配置为False时,此值需要与 FixedLossScaleManager 中的 loss_scale 相同。有关更多详细信息,请参阅 mindspore.FixedLossScaleManager。默认值:1.0。

输入:

  • gradients (tuple[Tensor]) - params 的梯度,shape与 params 相同。

输出:

Tensor[bool],值为True。

异常:

  • TypeError - learning_rate 不是int、float、Tensor、Iterable或LearningRateSchedule。

  • TypeError - parameters 的元素不是Parameter或字典。

  • TypeError - beta1beta2epsloss_scale 不是float。

  • TypeError - weight_decay 不是float或int。

  • TypeError - use_lockinguse_nesterov 不是bool。

  • ValueError - loss_scaleeps 不大于0。

  • ValueError - beta1beta2 不在(0.0,1.0)范围内。

  • ValueError - weight_decay 小于0。

支持平台:

Ascend GPU CPU

样例:

>>> from mindspore import nn, Model
>>>
>>> net = Net()
>>> #1) All parameters use the same learning rate and weight decay
>>> optim = nn.AdamOffload(params=net.trainable_params())
>>>
>>> #2) Use parameter groups and set different values
>>> conv_params = list(filter(lambda x: 'conv' in x.name, net.trainable_params()))
>>> no_conv_params = list(filter(lambda x: 'conv' not in x.name, net.trainable_params()))
>>> group_params = [{'params': conv_params, 'weight_decay': 0.01},
...                 {'params': no_conv_params, 'lr': 0.01},
...                 {'order_params': net.trainable_params()}]
>>> optim = nn.AdamOffload(group_params, learning_rate=0.1, weight_decay=0.0)
>>> # The conv_params's parameters will use default learning rate of 0.1 and weight decay of 0.01.
>>> # The no_conv_params's parameters will use learning rate of 0.01 and default weight decay of 0.0.
>>> # The final parameters order in which the optimizer will be followed is the value of 'order_params'.
>>>
>>> loss = nn.SoftmaxCrossEntropyWithLogits()
>>> model = Model(net, loss_fn=loss, optimizer=optim)