mindspore.nn.Lamb
- class mindspore.nn.Lamb(params, learning_rate, beta1=0.9, beta2=0.999, eps=1e-06, weight_decay=0.0)[源代码]
LAMB(Layer-wise Adaptive Moments optimizer for Batching training,用于批训练的分层自适应矩优化器)算法的实现。
LAMB是一种采用分层自适应批优化技术的优化算法。详见论文 LARGE BATCH OPTIMIZATION FOR DEEP LEARNING: TRAINING BERT IN 76 MINUTES。
LAMB优化器旨在不降低精度的情况下增加训练batch size,支持自适应逐元素更新和精确的分层校正。
参数更新如下:
\[\begin{split}\begin{array}{l} &\newline &\hline \\ &\textbf{Parameters}: \: 1^{\text {st }}\text {moment vector} \: m , \: 2^{\text {nd}} \: \text{moment vector} \: v , \\ &\hspace{5mm}\text{learning rate } \left\{ \gamma_{t}\right\}_{t=1}^{T} , \: \text {exponential decay rates for the moment estimates} \: \beta_{1} \: \beta_{2} , \\ &\hspace{5mm}\text{scaling function } \phi \\ &\textbf{Init}: \boldsymbol{m}_{0} \leftarrow 0, \: \boldsymbol{v}_{0} \leftarrow 0 \\[-1.ex] &\newline &\hline \\ &\textbf{for} \text { t=1 to T } \textbf{do} \\ &\hspace{5mm}\text{Draw b samples } S_{t} \text{ from } \mathbb{P} \text{ . } \\ &\hspace{5mm}\text{Compute } g_{t}=\frac{1}{\left|\mathcal{S}_{t}\right|} \sum_{s_{t} \in \mathcal{S}_{t}} \nabla \ell\left(x_{t}, s_{t}\right) . \\ &\hspace{5mm}\boldsymbol{m}_{t} \leftarrow \beta_{1} \boldsymbol{m}_{t-1}+\left(1-\beta_{1}\right) \boldsymbol{g}_{t} \\ &\hspace{5mm}\boldsymbol{v}_{t} \leftarrow \beta_{2} \boldsymbol{v}_{t-1}+\left(1-\beta_{2}\right) \boldsymbol{g}_{t}^{2} \\ &\hspace{5mm}\hat{\boldsymbol{m}}_{t} \leftarrow \boldsymbol{m}_{t} /\left(1-\beta_{1}^{t}\right) \\ &\hspace{5mm}\hat{\boldsymbol{v}}_{t} \leftarrow \boldsymbol{v}_{t} /\left(1-\beta_{2}^{t}\right) \\ &\hspace{5mm}\text{Compute ratio } \boldsymbol{r}_{t}=\hat{\boldsymbol{m}}_{t} /(\sqrt{\hat{\boldsymbol{v}}_{t}}+\epsilon) \\ &\hspace{5mm}\boldsymbol{w}_{t+1}^{(i)}=\boldsymbol{w}_{t}^{(i)}- \gamma_{t} \frac{\boldsymbol{\phi}\left(\left\|\boldsymbol{w}_{t}^{(i)}\right\|\right)} {\left\|\boldsymbol{r}_{t}^{(i)}+\lambda \boldsymbol{w}_{t}^{(i)}\right\|}\left(\boldsymbol{r}_{t}^{(i)}+ \lambda \boldsymbol{w}_{t}^{(i)}\right) \\ &\textbf{end for} \\[-1.ex] &\newline &\hline \\[-1.ex] &\textbf{return} \: \boldsymbol{w}_{t+1}\\[-1.ex] &\newline &\hline \\[-1.ex] \end{array}\end{split}\]其中, \(m\) 代表第一个动量矩阵 moment1 ,\(v\) 代表第二个动量矩阵 moment2 ,\(g\) 代表梯度 gradients ,\(\gamma\) 代表学习率 learning_rate,\(\beta_1, \beta_2\) 代表衰减速率 beta1 和 beta2 ,\(t\) 代表当前step,\(beta_1^t\) 和 \(beta_2^t\) 代表 beta1 和 beta2 的t次方 , \(w\) 代表 params , \(\epsilon\) 代表 eps, \(\lambda\) 表示LAMB权重衰减率。
说明
在参数未分组时,优化器配置的 weight_decay 应用于名称不含”beta”或”gamma”的网络参数。
用户可以分组调整权重衰减策略。分组时,每组网络参数均可配置 weight_decay 。若未配置,则该组网络参数使用优化器中配置的 weight_decay 。
优化器和混合精度之间通常没有联系。但是,当使用 FixedLossScaleManager 且 FixedLossScaleManager 中的 drop_overflow_update 设置为False时,优化器需要设置’loss_scale’。
由于此优化器没有 loss_scale 的参数,因此需要通过其他方式处理 loss_scale 。
如何正确处理 loss_scale 详见 LossScale。
警告
Lamb优化器的更新过程不是完全element-wise的,分布式并行下对权重进行切分可能会影响更新结果。
- 参数:
params (Union[list[Parameter], list[dict]]) - 必须是 Parameter 组成的列表或字典组成的列表。当列表元素是字典时,字典的键可以是”params”、”lr”、”weight_decay”、”grad_centralization”和”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值。
grad_centralization - 可选。如果键中存在”grad_centralization”,则使用对应的值,该值必须为布尔类型。如果没有,则认为 grad_centralization 为False。该参数仅适用于卷积层。
order_params - 可选。值的顺序是参数更新的顺序。当使用参数分组功能时,通常使用该配置项保持 parameters 的顺序以提升性能。如果键中存在”order_params”,则会忽略该组配置中的其他键。”order_params”中的参数必须在某一组 params 参数中。
learning_rate (Union[float, int, Tensor, Iterable, LearningRateSchedule]) -
float - 固定的学习率。必须大于等于零。
int - 固定的学习率。必须大于等于零。整数类型会被转换为浮点数。
Tensor - 可以是标量或一维向量。标量是固定的学习率。一维向量是动态的学习率,第i步将取向量中第i个值作为学习率。
Iterable - 动态的学习率。第i步将取迭代器第i个值作为学习率。
LearningRateSchedule - 动态的学习率。在训练过程中,优化器将使用步数(step)作为输入,调用 LearningRateSchedule 实例来计算当前学习率。
beta1 (float) - 第一矩的指数衰减率。参数范围(0.0,1.0)。默认值:
0.9
。beta2 (float) - 第二矩的指数衰减率。参数范围(0.0,1.0)。默认值:
0.999
。eps (float) - 将添加到分母中,以提高数值稳定性。必须大于0。默认值:
1e-6
。weight_decay (Union[float, int, Cell]) - 权重衰减(L2 penalty)。默认值:
0.0
。float: 固定值,必须大于或者等于0。
int: 固定值,必须大于或者等于0,会被转换成float。
Cell: 动态weight decay。在训练过程中,优化器会使用步数(step)作为输入,调用该Cell实例来计算当前weight decay值。
- 输入:
gradients (tuple[Tensor]) - params 的梯度,shape与 params 相同。
- 输出:
tuple[bool],所有元素都为True。
- 异常:
TypeError - learning_rate 不是int、float、Tensor、Iterable或LearningRateSchedule。
TypeError - parameters 的元素不是Parameter或dict。
TypeError - beta1、beta2 或 eps 不是float。
TypeError - weight_decay 不是float或int。
ValueError - eps 小于等于0。
ValueError - beta1、beta2 不在(0.0,1.0)范围内。
ValueError - weight_decay 小于0。
- 支持平台:
Ascend
GPU
样例:
>>> import mindspore as ms >>> from mindspore import nn >>> >>> # Define the network structure of LeNet5. Refer to >>> # https://gitee.com/mindspore/docs/blob/r2.3.0rc2/docs/mindspore/code/lenet.py >>> net = LeNet5() >>> #1) All parameters use the same learning rate and weight decay >>> optim = nn.Lamb(params=net.trainable_params(), learning_rate=0.1) >>> >>> #2) Use parameter groups and set different values >>> poly_decay_lr = nn.PolynomialDecayLR(learning_rate=0.1, end_learning_rate=0.01, ... decay_steps=4, power = 0.5) >>> 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, 'grad_centralization':True}, ... {'params': no_conv_params, 'lr': poly_decay_lr}, ... {'order_params': net.trainable_params(0.01)}] >>> optim = nn.Lamb(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 and grad >>> # centralization of True. >>> # The no_conv_params's parameters will use dynamic learning rate of poly decay learning rate and default >>> # weight decay of 0.0 and grad centralization of False. >>> # The final parameters order in which the optimizer will be followed is the value of 'order_params'. >>> >>> loss = nn.SoftmaxCrossEntropyWithLogits() >>> model = ms.train.Model(net, loss_fn=loss, optimizer=optim)