时间溢出的地方
0%

PyTorch 中的优化器和学习率调整机制

发表于 分类于 阅读次数： Valine：
本文字数： 23k 阅读时长 ≈ 21 分钟

1. 优化器的通用用法

首先引入包：

1
from torch import optim

使用例子：

1
2
optimizer = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
optimizer = optim.Adam([var1, var2], lr=0.0001)

优化器类初始化的第一个参数需要是一个 iterable，比如 model.parameters()。列表也是 iterable，比如 [var1, var2]。这个 iterable 里的元素应当是 Variabel（也可以是 dict）。

官方文档里提到一句话：

If you need to move a model to GPU via .cuda(), please do so before constructing optimizers for it. Parameters of a model after .cuda() will be different objects with those before the call.

也就是说，将模型加载到 GPU 之后再定义优化器（如果这个优化器更新的是这个模型的参数的话）。

1.1. 模型的不同部分使用不同的优化器超参数

实现这样的功能，传入优化器初始化的第一个参数就要是一个元素为 dict 的 iterable，并且 dict 应该有一个 key 是 "params", 其 item 是相应的模块的 parameters。其他的 key 应该是这个优化器初始化的参数名（如 lr, weight_decay 等）。举个例子：

1
2
3
4
optim.SGD([
                {'params': model.base.parameters()},
                {'params': model.classifier.parameters(), 'lr': 1e-3}
            ], lr=1e-2, momentum=0.9)

其中，model.base 部分的学习率为 1e-2，model.classifier 部分的学习率为 1e-3，二者的 momentum 都是 0.9。

1.2. 冻结模型部分层的参数

首先把需要冻结的层的参数设置为不需要求导：

1
2
for para in net.frozen_layer.parameters():
	para.requires_grad = False

然后在设置优化器的时候做一下参数筛选：

1
optimizer = optim.SGD(filter(lambda p: p.requires_grad, model.parameters()), lr=0.01, momentum=0.9)

1.3. 训练时使用优化器的步骤

这样做：

1
2
3
4
5
6
for input, target in dataset:
    optimizer.zero_grad()
    output = model(input)
    loss = loss_fn(output, target)
    loss.backward()
    optimizer.step()

2. 基类优化器 Optimizer

PyTorch 中的优化器都继承自基类 class optim.Optimizer，Optimizer 类的 __init__ 方法定义为：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
# __init__ method of optim.Optimizer
def __init__(self, params, defaults):
    torch._C._log_api_usage_once("python.optimizer")
    self.defaults = defaults

    self._hook_for_profile()

    if isinstance(params, torch.Tensor):
        raise TypeError("params argument given to the optimizer should be "
                        "an iterable of Tensors or dicts, but got " +
                        torch.typename(params))

    self.state = defaultdict(dict)
    self.param_groups = []

    param_groups = list(params)
    if len(param_groups) == 0:
        raise ValueError("optimizer got an empty parameter list")
    if not isinstance(param_groups[0], dict):
        param_groups = [{'params': param_groups}]

    for param_group in param_groups:
        self.add_param_group(param_group)

它的两个参数，其中的 params 就是前面说的需要传递进来的 iterable，defaluts 是一个字典，包含了优化器的默认参数（lr, weight_decay 等）。

文档里的说法是：

params: an iterable of torch.Tensor s or dict s. Specifies what Tensors should be optimized.

defaults: a dict containing default values of optimization options (used when a parameter group doesn’t specify them).

这个 defaults 在 Optimizer 类中并没有默认定义，但在它的子类中，即具体的优化器里，会给它一个定义，并传递给 Optimizer 类。举个例子，下面是 SGD 优化器类的 __init__ 函数（optim.SGD），可以看到其中定义了一个 defaults 变量，将该函数的默认参数打包成一个字典，并传递给了父类的 __init__ 函数。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
# __init__ method of optim.SGD
def __init__(self, params, lr=required, momentum=0, dampening=0,
                 weight_decay=0, nesterov=False):
        if lr is not required and lr < 0.0:
            raise ValueError("Invalid learning rate: {}".format(lr))
        if momentum < 0.0:
            raise ValueError("Invalid momentum value: {}".format(momentum))
        if weight_decay < 0.0:
            raise ValueError("Invalid weight_decay value: {}".format(weight_decay))

        defaults = dict(lr=lr, momentum=momentum, dampening=dampening,
                        weight_decay=weight_decay, nesterov=nesterov)
        if nesterov and (momentum <= 0 or dampening != 0):
            raise ValueError("Nesterov momentum requires a momentum and zero dampening")
        super(SGD, self).__init__(params, defaults)

再回来看 Optimizer 类的 __init__ 函数，其中的第 14 行定义了 self.param_groups = []，将会存储该优化器将要优化的参数。

第 16 行 param_groups = list(params) 将参数 params 变成列表，前面提到，params 作为一个 iterable，有三种形式：

  1. model.parameters(): 这里的 model 是 nn.Module 对象，它的方法 parameters() 返回一个 generator 对象，这当然是一个 iterator，把 model.parameters() 列表化，得到一个列表，其中的元素是 torch.nn.parameter.Parameter 对象，也就是模型层的权重了。
  2. [var1, var2]: 这本来就是个列表，用 list() 方法后还是原来的列表。其中的 var1, var2 也是 torch.nn.parameter.Parameter 对象。
  3. [{'params': model.base.parameters()}, {'params': model.classifier.parameters(), 'lr': 1e-3}]: 这是一个列表，其中的元素是字典类型。列表用 list() 方法还是原来的列表。

这样第 16 行之后 param_groups 作为一个列表，其内部元素的类型有两种：(a) torch.nn.parameter.Parameter 对象；(b) 字典对象，字典里至少有一个 key 为 params 的 item，可能还有其他 key（lr, weight_decay 等）。

第 19 和 20 行做的事情，是让 param_groups 的元素类型必须为字典，如果你不是字典，我就帮你变成字典。如果param_groups 的元素类型为 torch.nn.parameter.Parameter 对象，那么就把 param_groups 这个列表本身作为 key 为 params 的 item 的 value，包装成一个字典，再赋值给 param_groups。

第 22 和 23 行函数 add_param_group 的主要作用是将 param_groups 的东西放进 self.param_groups 中，注意 param_groups 和 self.param_groups 是两个变量。add_param_group 的代码如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
# add_param_group method of optim.Optimizer
def add_param_group(self, param_group):
    r"""Add a param group to the :class:`Optimizer` s `param_groups`.
    This can be useful when fine tuning a pre-trained network as frozen layers can be made
    trainable and added to the :class:`Optimizer` as training progresses.
    Args:
        param_group (dict): Specifies what Tensors should be optimized along with group
        specific optimization options.
    """
    assert isinstance(param_group, dict), "param group must be a dict"

    params = param_group['params']
    if isinstance(params, torch.Tensor):
        param_group['params'] = [params]
    elif isinstance(params, set):
        raise TypeError('optimizer parameters need to be organized in ordered collections, but '
                        'the ordering of tensors in sets will change between runs. Please use a list instead.')
    else:
        param_group['params'] = list(params)

    for param in param_group['params']:
        if not isinstance(param, torch.Tensor):
            raise TypeError("optimizer can only optimize Tensors, "
                            "but one of the params is " + torch.typename(param))
        if not param.is_leaf:
            raise ValueError("can't optimize a non-leaf Tensor")

    for name, default in self.defaults.items():
        if default is required and name not in param_group:
            raise ValueError("parameter group didn't specify a value of required optimization parameter " +
                             name)
        else:
            param_group.setdefault(name, default)

    params = param_group['params']
    if len(params) != len(set(params)):
        warnings.warn("optimizer contains a parameter group with duplicate parameters; "
                      "in future, this will cause an error; "
                      "see github.com/pytorch/pytorch/issues/40967 for more information", stacklevel=3)

    param_set = set()
    for group in self.param_groups:
        param_set.update(set(group['params']))

    if not param_set.isdisjoint(set(param_group['params'])):
        raise ValueError("some parameters appear in more than one parameter group")

    self.param_groups.append(param_group)

第 12 行 params = param_group['params']，从前面的分析知道，param_group['params'] 要么是一个列表，要么是一个 iterable 对象（比如 model.base.parameters()），13 - 19 行做的事情，就是不管 param_group['params'] 是什么，都要成为一个列表，并且列表的元素是 torch.nn.Parameter 对象，其实也就是 Tensor。

第 21 - 26 行对 param_group['params'] 里的元素做检查。如果不是叶子 Tensor，就报错。（PyTorch 不保存非叶子节点的 Tensor 的梯度）。

第 28 - 33 行补充 param_group 里的其他键值对，比如 lr, weight_decay 等，如果这些值已经有了，那就不添加，如果没有，就用 self.defaults 里的值。这一点是为了实现不同层使用不同参数的功能。

之后对self.param_groups和param_group中的元素进行判断，确保没有重复的参数。最后将字典**param_group**放进列表**self.param_groups**中

Optimizer 的 zero_grad 函数是将所有参数的梯度置为零，代码如下。其中 detach_()的作用是”Detaches the Tensor from the graph that created it, making it a leaf.”

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
# zero_grad method of optim.Optimizer
def zero_grad(self, set_to_none: bool = False):
    r"""Sets the gradients of all optimized :class:`torch.Tensor` s to zero.
    Args:
        set_to_none (bool): instead of setting to zero, set the grads to None.
            This will in general have lower memory footprint, and can modestly improve performance.
            However, it changes certain behaviors. For example:
            1. When the user tries to access a gradient and perform manual ops on it,
            a None attribute or a Tensor full of 0s will behave differently.
            2. If the user requests ``zero_grad(set_to_none=True)`` followed by a backward pass, ``.grad``\ s
            are guaranteed to be None for params that did not receive a gradient.
            3. ``torch.optim`` optimizers have a different behavior if the gradient is 0 or None
            (in one case it does the step with a gradient of 0 and in the other it skips
            the step altogether).
    """
    if not hasattr(self, "_zero_grad_profile_name"):
        self._hook_for_profile()
    with torch.autograd.profiler.record_function(self._zero_grad_profile_name):
        for group in self.param_groups:
            for p in group['params']:
                if p.grad is not None:
                    if set_to_none:
                        p.grad = None
                    else:
                        if p.grad.grad_fn is not None:
                            p.grad.detach_()
                        else:
                            p.grad.requires_grad_(False)
                        p.grad.zero_()

Optimizer 更新参数主要是靠 step 函数，在父类 Optimizer 的 step 函数中只有一行代码 raise NotImplementedError ,这意味着每个子类都必须实现自己的 step 函数，这个原因很显然，因为不同的优化器更新参数的方式是不同的。后面将会对几个不同的优化器的 step 函数做一下分析。

另外还有两个比较常用的函数：state_dict() 和 load_staet_dict()。

state_dict() 函数返回一个字典，大致如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
# state_dict method of optim.Optimizer
def state_dict(self):
    r"""Returns the state of the optimizer as a :class:`dict`.
    It contains two entries:
    * state - a dict holding current optimization state. Its content
        differs between optimizer classes.
    * param_groups - a list containing all parameter groups where each
        parameter group is a dict
    """
    # codes here
    return {
        'state': packed_state,
        'param_groups': param_groups,
    }

保存模型的时候经常也会保存优化器，这时就需要用到 state_dict()：

1
torch.save(optimizer.state_dict(), "optimizer.pt")

load_state_dict() 函数加载保存的优化器状态：

1
2
optimizer_state_dict = torch.load("optimizer.pt")
optimizer.load_state_dict(optimizer)

3. 一些常用的优化器

3.1. SGD 优化器

类签名：

1
2
3
4
5
6
torch.optim.SGD(params, 
                lr=<required parameter>, 
                momentum=0, 
                dampening=0, 
                weight_decay=0, 
                nesterov=False)

参数含义：

  • params (iterable) – iterable of parameters to optimize or dicts defining parameter groups
  • lr (float) – learning rate
  • momentum (float, optional) – momentum factor (default: 0)
  • weight_decay (float, optional) – weight decay (L2 penalty) (default: 0)
  • dampening (float, optional) – dampening for momentum (default: 0)
  • nesterov (bool, optional) – enables Nesterov momentum (default: False)

计算流程可以参考：

optim.SGD 的代码里主要定义的就是 step() 函数，其主要步骤的实现函数 F.sgd 代码为：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
def sgd(params: List[Tensor],
        d_p_list: List[Tensor],
        momentum_buffer_list: List[Optional[Tensor]],
        *,
        weight_decay: float,
        momentum: float,
        lr: float,
        dampening: float,
        nesterov: bool):
    r"""Functional API that performs SGD algorithm computation.
    See :class:`~torch.optim.SGD` for details.
    """

    for i, param in enumerate(params):

        d_p = d_p_list[i]
        if weight_decay != 0:
            d_p = d_p.add(param, alpha=weight_decay)

        if momentum != 0:
            buf = momentum_buffer_list[i]

            if buf is None:
                buf = torch.clone(d_p).detach()
                momentum_buffer_list[i] = buf
            else:
                buf.mul_(momentum).add_(d_p, alpha=1 - dampening)

            if nesterov:
                d_p = d_p.add(buf, alpha=momentum)
            else:
                d_p = buf

        param.add_(d_p, alpha=-lr)

3.2. Adam 优化器

类签名：

1
2
3
4
5
6
torch.optim.Adam(params, 
                 lr=0.001, 
                 betas=(0.9, 0.999), 
                 eps=1e-08, 
                 weight_decay=0, 
                 amsgrad=False)

参数含义：

  • params (iterable) – iterable of parameters to optimize or dicts defining parameter groups
  • lr (float, optional) – learning rate (default: 1e-3)
  • betas (Tuple**[float, float]**, optional) – coefficients used for computing running averages of gradient and its square (default: (0.9, 0.999))
  • eps (float, optional) – term added to the denominator to improve numerical stability (default: 1e-8)
  • weight_decay (float, optional) – weight decay (L2 penalty) (default: 0)
  • amsgrad (boolean**, optional) – whether to use the AMSGrad variant of this algorithm from the paper On the Convergence of Adam and Beyond (default: False)

计算流程为：

从这个计算流程可以看出，PyTorch 实现里的 Adam 优化器是用 L2 正则化来代替 weight decay 的。

optim.Adam 里具体的函数 F.adam 为：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
def adam(params: List[Tensor],
         grads: List[Tensor],
         exp_avgs: List[Tensor],
         exp_avg_sqs: List[Tensor],
         max_exp_avg_sqs: List[Tensor],
         state_steps: List[int],
         *,
         amsgrad: bool,
         beta1: float,
         beta2: float,
         lr: float,
         weight_decay: float,
         eps: float):
    r"""Functional API that performs Adam algorithm computation.
    See :class:`~torch.optim.Adam` for details.
    """

    for i, param in enumerate(params):

        grad = grads[i]
        exp_avg = exp_avgs[i]
        exp_avg_sq = exp_avg_sqs[i]
        step = state_steps[i]

        bias_correction1 = 1 - beta1 ** step
        bias_correction2 = 1 - beta2 ** step

        if weight_decay != 0:
            grad = grad.add(param, alpha=weight_decay)

        # Decay the first and second moment running average coefficient
        exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
        exp_avg_sq.mul_(beta2).addcmul_(grad, grad.conj(), value=1 - beta2)
        if amsgrad:
            # Maintains the maximum of all 2nd moment running avg. till now
            torch.maximum(max_exp_avg_sqs[i], exp_avg_sq, out=max_exp_avg_sqs[i])
            # Use the max. for normalizing running avg. of gradient
            denom = (max_exp_avg_sqs[i].sqrt() / math.sqrt(bias_correction2)).add_(eps)
        else:
            denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)

        step_size = lr / bias_correction1

        param.addcdiv_(exp_avg, denom, value=-step_size)

其中的 step 变量应该相当于公式里的 t。

3.3. AdamW 优化器

类签名：

1
2
3
4
5
6
torch.optim.AdamW(params, 
                  lr=0.001, 
                  betas=(0.9, 0.999), 
                  eps=1e-08, 
                  weight_decay=0.01, 
                  amsgrad=False)

参数含义：

  • params (iterable) – iterable of parameters to optimize or dicts defining parameter groups
  • lr (float, optional) – learning rate (default: 1e-3)
  • betas (Tuple**[float, float]**, optional) – coefficients used for computing running averages of gradient and its square (default: (0.9, 0.999))
  • eps (float, optional) – term added to the denominator to improve numerical stability (default: 1e-8)
  • weight_decay (float, optional) – weight decay coefficient (default: 1e-2)
  • amsgrad (boolean**, optional) – whether to use the AMSGrad variant of this algorithm from the paper On the Convergence of Adam and Beyond (default: False)]

计算流程为：

从这个流程来看，没有用 L2 正则化来代替 weight decay，并且上来就做了 weight decay，原文的算法流程里是在最后做 weight decay 的，不过这个顺序没有什么影响。

optim.Adamw 里具体的函数 F.adamw 为：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
def adamw(params: List[Tensor],
          grads: List[Tensor],
          exp_avgs: List[Tensor],
          exp_avg_sqs: List[Tensor],
          max_exp_avg_sqs: List[Tensor],
          state_steps: List[int],
          *,
          amsgrad: bool,
          beta1: float,
          beta2: float,
          lr: float,
          weight_decay: float,
          eps: float):
    r"""Functional API that performs AdamW algorithm computation.
    See :class:`~torch.optim.AdamW` for details.
    """
    for i, param in enumerate(params):
        grad = grads[i]
        exp_avg = exp_avgs[i]
        exp_avg_sq = exp_avg_sqs[i]
        step = state_steps[i]

        # Perform stepweight decay
        param.mul_(1 - lr * weight_decay)

        bias_correction1 = 1 - beta1 ** step
        bias_correction2 = 1 - beta2 ** step

        # Decay the first and second moment running average coefficient
        exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
        exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2)
        if amsgrad:
            # Maintains the maximum of all 2nd moment running avg. till now
            torch.maximum(max_exp_avg_sqs[i], exp_avg_sq, out=max_exp_avg_sqs[i])
            # Use the max. for normalizing running avg. of gradient
            denom = (max_exp_avg_sqs[i].sqrt() / math.sqrt(bias_correction2)).add_(eps)
        else:
            denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps)

        step_size = lr / bias_correction1

        param.addcdiv_(exp_avg, denom, value=-step_size)

4. 调整学习率

PyTorch 调整学习率的类都在 torch.optim.lr_scheduler 中，如 torch.optim.lr_scheduler.ReduceLROnPlateau。

4.1. ReduceLROnPlateau

类签名：

1
2
3
4
5
6
7
8
9
10
torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, 
                                           mode='min', 
                                           factor=0.1, 
                                           patience=10, 
                                           threshold=0.0001, 
                                           threshold_mode='rel', 
                                           cooldown=0, 
                                           min_lr=0, 
                                           eps=1e-08, 
                                           verbose=False)

各个参数的含义如下：

  • optimizer (Optimizer) – Wrapped optimizer.
  • mode (str) – One of min, max. In min mode, lr will be reduced when the quantity monitored has stopped decreasing; in max mode it will be reduced when the quantity monitored has stopped increasing. Default: ‘min’.
  • factor (float) – Factor by which the learning rate will be reduced. new_lr = lr * factor. Default: 0.1.
  • patience (int) – Number of epochs with no improvement after which learning rate will be reduced. For example, if patience = 2, then we will ignore the first 2 epochs with no improvement, and will only decrease the LR after the 3rd epoch if the loss still hasn’t improved then. Default: 10.
  • threshold (float) – Threshold for measuring the new optimum, to only focus on significant changes. Default: 1e-4.
  • threshold_mode (str) – One of rel, abs. In rel mode, dynamic_threshold = best ( 1 + threshold ) in ‘max’ mode or best ( 1 - threshold ) in min mode. In abs mode, dynamic_threshold = best + threshold in max mode or best - threshold in min mode. Default: ‘rel’. “rel” 是 relative 的意思，”abs” 是 absolute 的意思
  • cooldown (int) – Number of epochs to wait before resuming normal operation after lr has been reduced. Default: 0. 学习率调整后，在 cooldown 个 epoch 内，学习率不参与调整。即使在这期间指标 在 patience 个 epoch 之内没有提升，也不调整学习率
  • min_lr (float or list) – A scalar or a list of scalars. A lower bound on the learning rate of all param groups or each group respectively. Default: 0.
  • eps (float) – Minimal decay applied to lr. If the difference between new and old lr is smaller than eps, the update is ignored. Default: 1e-8.
  • verbose (bool) – If True, prints a message to stdout for each update. Default: False.

这个类是直接对 optimizer 的 param_groups 变量的 lr 键进行调整的，可以从源码（torch.optim.lr_scheduler.ReduceLROnPlateau._reduce_lr%3A)）得知：

1
2
3
4
5
6
7
8
9
def _reduce_lr(self, epoch):
    for i, param_group in enumerate(self.optimizer.param_groups):
        old_lr = float(param_group['lr'])
        new_lr = max(old_lr * self.factor, self.min_lrs[i])
        if old_lr - new_lr > self.eps:
            param_group['lr'] = new_lr
            if self.verbose:
                print('Epoch {:5d}: reducing learning rate'
                      ' of group {} to {:.4e}.'.format(epoch, i, new_lr))

使用方法一般如下：

1
2
3
4
5
6
7
optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
scheduler = ReduceLROnPlateau(optimizer, 'min')
for epoch in range(10):
    train(...)
    val_loss = validate(...)
    # Note that step should be called after validate()
    scheduler.step(val_loss)

4.2. CosineAnnealingLR

类签名：

1
2
3
4
5
torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, 
                                           T_max, 
                                           eta_min=0, 
                                           last_epoch=- 1, 
                                           verbose=False)

官方文档给了一个公式说明，但看得不是很明白：

其中的 $\eta_{max}$ 就是优化器的初始学习率，是整个学习率调整过程中余弦值的峰值，$\eta_{min}$ 是最小学习率，手动设置，默认为 0，是余弦值的谷底值。$T_{max}$ 其实就是余弦函数的半周期值。该调整方法对学习率的改变可以做个图看看，代码为：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
from torch import optim
from torch.optim import lr_scheduler
import matplotlib.pyplot as plt

lr_list = []
model = squeezenet  # 这里随便设置一个网络就行

LR = 0.01
optimizer = optim.Adam(model.parameters(),lr = LR)
scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max = 10)

epochs = 100

for epoch in range(epochs):
    scheduler.step()
    lr_list.append(optimizer.state_dict()['param_groups'][0]['lr'])


# plt.xticks(range(1, epochs+1))  # 用于调整刻度值
plt.plot(range(1, epochs+1), lr_list, color='r')

plt.savefig("cosine_annealing_lr.png", dpi=100)

下面是学习率变化曲线，可以看出半周期是 10。

4.3. CosineAnnealingWarmRestarts

带热启动的余弦退火学习率调整机制，参见文章 SGDR: Stochastic Gradient Descent with Warm Restarts。这个机制，是学习率按照余弦函数下降，降到最低点时直接一步回到最高点，所以叫「热重启」。

类签名：

1
2
3
4
5
6
torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, 
                                                     T_0, 
                                                     T_mult=1, 
                                                     eta_min=0, 
                                                     last_epoch=- 1, 
                                                     verbose=False)

官方文档的说明如下：

这里的 $\eta_{max}$ 和 $\eta_{min}$ 的含义和 CosineAnnealingLR 里的含义是一样的。

$T_0$ 的含义是学习率第一次回到初始值的 epoch 位置。如果只说到这里，其实这个 $T_0$ 和 CosineAnnealingLR 里的 $T_{max}$ 的含义是一样的，都是余弦函数的半周期。但之所以这里用的下标是 0 而不是 $max$，并且也强调了是「第一次」回到初始值的 epoch 位置，是因为还要考虑 T_mult 这个参数。

T_mult 参数控制了学习率变化的速度：

  • 如果 T_mult=1，那么 $T_0$ 就是余弦函数的半周期（在热重启里，其实只有半个周期，不存在整个周期），即学习率会在 $T_0$, $2T_0$, $3T_0$, $\dots$ 处回到最大值（初始学习率）
  • 如果 T_mult>1，则学习率会在 $T_0$, $(1 + T_{mult})T_0$, $(1 + T_{mult} + T_{mult}^2)T_0$, $\dots$, $(1 + T_{mult} + T_{mult}^2 + \cdots + T_{mult}^i)T_0$ 处回到最大值
  • 如果 T_mult<1，会报错

对比一下二者的学习率变化曲线，代码为：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
lr_list = []
model = squeezenet

LR = 0.01
optimizer = optim.Adam(model.parameters(),lr = LR)
scheduler = lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0 = 10, T_mult=1) # T_mult may change to 2

epochs = 100

for epoch in range(epochs):
    scheduler.step()
    lr_list.append(optimizer.state_dict()['param_groups'][0]['lr'])

# plt.xticks(range(1, epochs+1))
plt.plot(range(1, epochs+1), lr_list, color='r')

CosineAnnealingWR with T_0 = 10 and T_mult = 1

CosineAnnealingWR with T_0 = 10 and T_mult = 2

可以看到，当 T_mult=1 时，学习率都是在第 10, 20, 30, … 个 epoch 的时候回到最大值。当 T_mult=2 时，学习率在第 10, 30, 70, … 个 epoch 的时候回到最大值。

在调节参数的时候，一定要根据自己总的epoch合理的设置参数，不然很可能达不到预期的效果,经过我自己的试验发现，如果是用那种等间隔的退火策略（CosineAnnealingLR 和 T*mult=1 的 CosineAnnealingWarmRestarts），验证准确率总是会在学习率的最低点达到一个很好的效果，而随着学习率回升，验证精度会有所下降.所以为了能最终得到一个更好的收敛点，设置 T_mult>1 是很有必要的，这样到了训练后期，学习率不会再有一个回升的过程,而且一直下降直到训练结束。

引用自：pytorch的余弦退火学习率

这个优化器使用时，可以直接使用 step() 更新，如下：

1
2
3
4
5
6
7
optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
scheduler = CosineAnnealingWarmRestarts(optimizer, T_0 = 10, T_mult=1)
for epoch in range(10):
    train(...)
    val_loss = validate(...)
    # Note that step should be called after validate()
    scheduler.step()

这时是以 epoch 为单位更新，即在一个 epoch 内的所有 batch，都用同一个学习率。

还可以以 batch 为单位更新（这个使用方法来源于官方文档），如下：

1
2
3
4
5
6
7
8
9
10
11
scheduler = CosineAnnealingWarmRestarts(optimizer, T_0, T_mult)
iters = len(dataloader)
for epoch in range(20):
    for i, sample in enumerate(dataloader):
        inputs, labels = sample['inputs'], sample['labels']
        optimizer.zero_grad()
        outputs = net(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        scheduler.step(epoch + i / iters)  # 这一行

一点疑惑：

关于这个以 batch 为单位，很自然的一种想法是我还是用 step() 更新，但我是每个 batch 的时候都使用一次 scheduler.step()，那么这个和使用 scheduler.step(epoch + i / iters) 会不会是一样的呢？我做了一下对比。

首先是每个 batch 使用 scheduler.step() 函数，把这种情况叫做「step 无参数」：

n
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
lr_list = []
model = squeezenet

LR = 0.01
optimizer = optim.Adam(model.parameters(),lr = LR)
scheduler = lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0 = 10, T_mult=1)

epochs = 20
batches = 5

for epoch in range(epochs):
    for batch in range(batches):
        scheduler.step()
        lr_list.append(optimizer.state_dict()['param_groups'][0]['lr'])


plt.plot(range(1, epochs*batches+1), lr_list, color='r')

它的学习率变化曲线如下：

另一种是每个 batch 使用 scheduler.step(epoch + i / iters) 函数，把这种情况叫做「step 有参数」：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
lr_list = []
model = squeezenet

LR = 0.01
optimizer = optim.Adam(model.parameters(),lr = LR)
scheduler = lr_scheduler.CosineAnnealingWarmRestarts(optimizer, T_0 = 10, T_mult=1)

epochs = 20
batches = 5

for epoch in range(epochs):
    for batch in range(batches):
        scheduler.step(epoch + batch / batches)
        lr_list.append(optimizer.state_dict()['param_groups'][0]['lr'])


plt.plot(range(1, epochs*batches+1), lr_list, color='r')

它的学习率变化曲线如下所示：

可以看出两种方法对学习率调整的情况实际上是不一样的。分析这两种不同时，我们先把每次调用 scheduler 的 step 函数这个操作叫做「步」。两种方法的 epoch 都设置为 20，每个 epoch 的 batch 数量都为 5。

对于 「step 无参数」的方法，其实是在每一个 batch 都让 scheduler 作用了一步，所以这种方法的学习率变化曲线跟设置 100 个 epoch，每个 epoch 更新一次的曲线是一样的。

对于「step 有参数」的方法，每个 batch 时 scheduler 作用的不是一步，而是 0.2 步（1 除以 batch 数量），这样需要作用 50 个 batch 步数才到达 $T_0=10$，也就形成了曲线里横坐标为 50 的时候学习率回到最高值。

上面的解释如果有点模糊，可以看一下 CosineAnnealingWarmRestarts 的 step 函数源代码：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
def step(self, epoch=None):
    if epoch is None and self.last_epoch < 0:
        epoch = 0

    if epoch is None:
        epoch = self.last_epoch + 1
        self.T_cur = self.T_cur + 1
        if self.T_cur >= self.T_i:
            self.T_cur = self.T_cur - self.T_i
            self.T_i = self.T_i * self.T_mult
    else:
        if epoch < 0:
            raise ValueError("Expected non-negative epoch, but got {}".format(epoch))
        if epoch >= self.T_0:
            if self.T_mult == 1:
                self.T_cur = epoch % self.T_0
            else:
                n = int(math.log((epoch / self.T_0 * (self.T_mult - 1) + 1), self.T_mult))
                self.T_cur = epoch - self.T_0 * (self.T_mult ** n - 1) / (self.T_mult - 1)
                self.T_i = self.T_0 * self.T_mult ** (n)
        else:
            self.T_i = self.T_0
            self.T_cur = epoch
    self.last_epoch = math.floor(epoch)

这个函数里有个参数 epoch，可以把这个 epoch 理解为上面所说的「步」。使用「step 无参数」方法时，每次调用 step，epoch 参数都会设置为 None，根据第 5 - 6 行，这时候会把 epoch 加 1，也就是说，这时候「步」会加 1（这时「步」一直会是整数）。而使用「step 有参数」方法时，每次调用 step，会执行第 12 - 23 行的代码，这时的「步」不会自动加 1（因为参数已经提供了 epoch 值，epoch 每次加 0.2）。

以上两种情况的不同，仍然觉得没有解释的好，只能根据代码大概意会，等以后再继续学习吧。

4.4. LambdaLR

类签名：

1
2
3
4
torch.optim.lr_scheduler.LambdaLR(optimizer, 
                                  lr_lambda, 
                                  last_epoch=- 1, 
                                  verbose=False)

自定义学习率变化，通过 lambda_lr，一个自定义函数来调整，lambda_lr 不是直接返回学习率，而是返回一个与学习率相乘的因子，来确定最终的实际学习率。lambda_lr 可以是一个函数，也可以是一个函数列表，用于对不同的 param_groups 实行不同的学习率调整机制。lambda_lr 的参数是一个表示 epoch 的单参数。

官方文档的参数说明：

  • optimizer (Optimizer) – Wrapped optimizer.
  • lr_lambda (function or list) – A function which computes a multiplicative factor given an integer parameter epoch, or a list of such functions, one for each group in optimizer.param_groups.
  • last_epoch (int) – The index of last epoch. Default: -1.
  • verbose (bool) – If True, prints a message to stdout for each update. Default: False.

官方示例：

1
2
3
4
5
6
7
8
# Assuming optimizer has two groups.
lambda1 = lambda epoch: epoch // 30
lambda2 = lambda epoch: 0.95 ** epoch
scheduler = LambdaLR(optimizer, lr_lambda=[lambda1, lambda2])
for epoch in range(100):
    train(...)
    validate(...)
    scheduler.step()

HuggingFace 的 pytorch_transformers 库里的学习率调整机制就是继承 LambdaLR 类的。

5. Reference

  1. PyTorch optim 官方文档
  2. 知乎文章：【PyTorch】优化器 torch.optim.Optimizer
  3. pytorch的余弦退火学习率