Pytorch 自定义LRScheduler

Pytorch 自定义 PolyScheduler

---

写在前面

最近有个自定义学习率优化策略的需求，但是网上搜索到的大部分都是Pytorch内部的LambdaLR.，虽然简单易用，但是无法满足需求。所以就研究了下Pytorch的源码，仿照源码写了一个PolyScheduler。有两个版本，一个版本是变化检测开源库中的版本：likyoo的变化检测库，另一个是本文中的版本，两个版本都是可行的，大家根据自己习惯去复制。

---

一、PolyScheduler代码用法

在代码开始之前，我们先看一下算法本身：
$$lr=base\_lr\times {\left(1-\frac{step}{total\_step}\right)}^{power}$$

这里的step可以是epoch，也可以是每个batch之后step一次。

每个batch之后step一次的用法：

data_loader = torch.utils.data.DataLoader(...)
steps_per_epoch = len(data_loader)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
scheduler = torch.optim.lr_scheduler.PolyScheduler(optimizer, min_lr=0.01, steps_per_epoch=steps_per_epoch , epochs=10)
for epoch in range(10):
	for i,batch in enumerate(enumerate):
		train_batch(...)
		optimizer.step()
		scheduler.step(epoch*steps_per_epoch + i)

每个epoch之后step的用法：

data_loader = torch.utils.data.DataLoader(...)
optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
scheduler = torch.optim.lr_scheduler.PolyScheduler(optimizer, min_lr=0.01,epochs=10)
for epoch in range(10):
	for i,batch in enumerate(enumerate):
		train_batch(...)
		optimizer.step()
	
	scheduler.step()

二、PolyScheduler源码

from torch.optim.lr_scheduler import _LRScheduler

class LinoPolyScheduler(_LRScheduler):
    r"""
    Args:
        optimizer (Optimizer): Wrapped optimizer.
        total_steps (int): The total number of steps in the cycle. Note that
            if a value is not provided here, then it must be inferred by providing
            a value for epochs and steps_per_epoch.
            Default: None
        epochs (int): The number of epochs to train for. This is used along
            with steps_per_epoch in order to infer the total number of steps in the cycle
            if a value for total_steps is not provided.
            Default: None
        steps_per_epoch (int): The number of steps per epoch to train for. This is
            used along with epochs in order to infer the total number of steps in the
            cycle if a value for total_steps is not provided.
            Default: None
        last_epoch (int): The index of the last batch. This parameter is used when
            resuming a training job. Since `step()` should be invoked after each
            batch instead of after each epoch, this number represents the total
            number of *batches* computed, not the total number of epochs computed.
            When last_epoch=-1, the schedule is started from the beginning.
            Default: -1
        verbose (bool): If ``True``, prints a message to stdout for
            each update. Default: ``False``.

    Example:
        >>> data_loader = torch.utils.data.DataLoader(...)
        >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
        >>> scheduler = torch.optim.lr_scheduler.PolyScheduler(optimizer, min_lr=0.01, steps_per_epoch=None, epochs=10)
        >>> for epoch in range(10):
        >>>     for batch in data_loader:
        >>>         train_batch(...)
        >>>     scheduler.step()


    .. _Super-Convergence\: Very Fast Training of Neural Networks Using Large Learning Rates:
        https://arxiv.org/abs/1708.07120
    """

    def __init__(self,
                 optimizer,
                 power=1.0,
                 total_steps=None,
                 epochs=None,
                 steps_per_epoch=None,
                 min_lr=0,
                 last_epoch=-1,
                 verbose=False):

        # Validate optimizer
        if not isinstance(optimizer, Optimizer):
            raise TypeError('{} is not an Optimizer'.format(
                type(optimizer).__name__))
        self.optimizer = optimizer
        # self.by_epoch = by_epoch
        self.epochs = epochs
        self.min_lr = min_lr
        self.power = power

        # check param
        param_dic = {'total_steps': total_steps, 'epochs': epochs, 'steps_per_epoch': steps_per_epoch}
        for k, v in param_dic.items():
            if v is not None:
                if v <= 0 or not isinstance(v, int):
                    raise ValueError("Expected positive integer {}, but got {}".format(k, v))

        # Validate total_steps
        if total_steps is not None:
            self.total_steps = total_steps
        elif epochs is not None and steps_per_epoch is None:
            self.total_steps = epochs
        elif epochs is not None and steps_per_epoch is not None:
            self.total_steps = epochs * steps_per_epoch
        else:
            raise ValueError("You must define either total_steps OR epochs OR (epochs AND steps_per_epoch)")

        super(LinoPolyScheduler, self).__init__(optimizer, last_epoch, verbose)

    def _format_param(self, name, optimizer, param):
        """Return correctly formatted lr/momentum for each param group."""
        if isinstance(param, (list, tuple)):
            if len(param) != len(optimizer.param_groups):
                raise ValueError("expected {} values for {}, got {}".format(
                    len(optimizer.param_groups), name, len(param)))
            return param
        else:
            return [param] * len(optimizer.param_groups)

    def get_lr(self):
        if not self._get_lr_called_within_step:
            warnings.warn("To get the last learning rate computed by the scheduler, "
                          "please use `get_last_lr()`.", UserWarning)

        step_num = self.last_epoch

        if step_num > self.total_steps:
            raise ValueError("Tried to step {} times. The specified number of total steps is {}"
                             .format(step_num + 1, self.total_steps))

        coeff = (1 - step_num / self.total_steps) ** self.power

        return [(base_lr - self.min_lr) * coeff + self.min_lr for base_lr in self.base_lrs]

---

三、如何在Pytorch中自定义学习率优化策略

学习率优化策略，顾名思义就是动态改变optimizer的学习率，那说到底是对optimizer做操作。在阅读Pytorch源码的过程中，学习率优化策略最基础的一个类便是_LRScheduler，我们自己去自定义也是需要继承自这个类的。

我们都知道，lr_scheduler除了最开始的实例化，另外一个步骤就是step()了，所以我们只需要看源码中step做了啥事情，就能搞懂如何自定义学习率优化策略了。

由于源代码较多，我们只看最关键的部分，上图中便是我从step()方法源码中截图出来两个关键的点，一个点是获取学习率，一个点是更新学习率。更新学习率我们没啥好看的，基本上都是一样的，最主要的便是获取学习率：get_lr()这个方法。在_LRScheduler类中并没有写此方法内容，就只有raise NotImplementedError，就意味着我们自定义的时候必须重写此方法。以咱们的PolyScheduler为例：

其中self.last_epoch是一个贯穿全局的变量，我们会在step(epoch=None)方法中传入当前的step，源码中会自动将self.last_epoch更新为我们传入的值。

我们便是以这种思路去自定义学习率优化策略的。

---

总结

关注公众号，跟Lino一起折腾吧！