mindspore.nn.StepLR

class mindspore.nn.StepLR(optimizer, step_size, gamma=0.5, last_epoch=- 1, verbose=False)[源代码]

step_size 个epoch按 gamma 衰减每个参数组的学习率。StepLR 对于学习率的衰减可能与外部对于学习率的改变同时发生。

警告

这是一个实验性的动态学习率接口,需要和 实验性优化器 下的接口配合使用。

参数:
  • optimizer (mindspore.nn.optim_ex.Optimizer) - 优化器实例。

  • step_size (int) - 学习率衰减的周期。

  • gamma (float,可选) - 学习率衰减的乘法因子。默认值: 0.1

  • last_epoch (int,可选) - epoch/step数。默认值: -1

  • verbose (bool,可选) - 是否打印学习率。默认值: False

支持平台:

Ascend GPU CPU

样例:

>>> import mindspore
>>> from mindspore import nn
>>> # Define the network structure of LeNet5. Refer to
>>> # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/lenet.py
>>> net = LeNet5()
>>> loss_fn = nn.SoftmaxCrossEntropyWithLogits(sparse=True)
>>> optimizer = nn.optim_ex.Adam(net.trainable_params(), lr=0.05)
>>> # Assuming optimizer uses lr = 0.05 for all groups
>>> # lr = 0.05     if epoch < 2
>>> # lr = 0.005    if 2 <= epoch < 4
>>> # lr = 0.0005   if 4 <= epoch < 6
>>> scheduler = nn.StepLR(optimizer, step_size=2, gamma=0.1)
>>> def forward_fn(data, label):
...     logits = net(data)
...     loss = loss_fn(logits, label)
...     return loss, logits
>>> grad_fn = mindspore.value_and_grad(forward_fn, None, optimizer.parameters, has_aux=True)
>>> def train_step(data, label):
...     (loss, _), grads = grad_fn(data, label)
...     optimizer(grads)
...     return loss
>>> for epoch in range(6):
...     # Create the dataset taking MNIST as an example. Refer to
...     # https://gitee.com/mindspore/docs/blob/r2.1/docs/mindspore/code/mnist.py
...     for data, label in create_dataset():
...         train_step(data, label)
...     scheduler.step()
...     current_lr = scheduler.get_last_lr()