[PyTorch] learning rate scheduler - StepLR에 대해

출처: https://github.com/pytorch/pytorch/blob/v1.0.1/torch/optim/lr_scheduler.py#L126

class StepLR(_LRScheduler):
    """Sets the learning rate of each parameter group to the initial lr
    decayed by gamma every step_size epochs. When last_epoch=-1, sets
    initial lr as lr.

    Args:
        optimizer (Optimizer): Wrapped optimizer.
        step_size (int): Period of learning rate decay.
        gamma (float): Multiplicative factor of learning rate decay.
            Default: 0.1.
        last_epoch (int): The index of last epoch. Default: -1.

    Example:
        >>> # Assuming optimizer uses lr = 0.05 for all groups
        >>> # lr = 0.05     if epoch < 30
        >>> # lr = 0.005    if 30 <= epoch < 60
        >>> # lr = 0.0005   if 60 <= epoch < 90
        >>> # ...
        >>> scheduler = StepLR(optimizer, step_size=30, gamma=0.1)
        >>> for epoch in range(100):
        >>>     scheduler.step()
        >>>     train(...)
        >>>     validate(...)
    """

    def __init__(self, optimizer, step_size, gamma=0.1, last_epoch=-1):
        self.step_size = step_size
        self.gamma = gamma
        super(StepLR, self).__init__(optimizer, last_epoch)

    def get_lr(self):
        return [base_lr * self.gamma ** (self.last_epoch // self.step_size)
                for base_lr in self.base_lrs]

 

각 매개변수 그룹의 학습률을 초기 lr로 설정하고, step_size를 에포크마다 감마에 의해 줄어듦드는 방식. 이런 걸 learning rate scheduler이라 함. 위 figure의 Step-based Decay를 보면 계단 형식으로 감소하는 것을 볼 수 있음

 

필요한 파라미터는 optimizer, step_size, gamma이다. 사용하는 옵티마이저 설정, 몇 에포크마다 감마만큼 줄어들게 할 것인지를 설정해주면 됨

 

def adjust_learning_rate(optimizer, iteration_count):
    """Imitating the original implementation"""
    lr = args.lr / (1.0 + args.lr_decay * iteration_count)
    for param_group in optimizer.param_groups:
        param_group['lr'] = lr

 

코드를 보다가 직접 구현한 코드를 보았음

 

이 코드 같은 경우엔 adjust_learning_rate 함수를 사용하여 최적화(optimizer) 객체의 학습률(learning rate)을 조정하는 데 사용하였고, iteration_count를 통해 학습이 진행됨에 따라 학습률을 감소시켰음. 

 

위의 StepLR과 차이를 살펴보면

1. StepLR은 일정한 간격으로 학습률을 감소 시킨다. 감소율에 대한 decay_factor 등을 지정할 수 있음. 대신 adjust_learning_rate는 간단한 선형적인 감소 방식으로 학습률을 조정함

2. 매개변수의 차이

3. StepLR은 PyTorch의 내장 기능으로 클래스를 초기화하고 옵티마이저의 step() 함수를 호출하여 학습률을 자동으로 업데이트 할 수 있음