chenxuanhong
135 lines
4.8 KiB
Python
135 lines
4.8 KiB
Python
#!/usr/bin/env python3
|
|
# -*- coding:utf-8 -*-
|
|
#############################################################
|
|
# File: learningrate_scheduler.py
|
|
# Created Date: Tuesday January 5th 2021
|
|
# Author: Chen Xuanhong
|
|
# Email: chenxuanhongzju@outlook.com
|
|
# Last Modified: Tuesday, 5th January 2021 2:04:00 pm
|
|
# Modified By: Chen Xuanhong
|
|
# Copyright (c) 2021 Shanghai Jiao Tong University
|
|
#############################################################
|
|
|
|
# Refer to basicSR https://github.com/xinntao/BasicSR
|
|
|
|
|
|
import math
|
|
from collections import Counter
|
|
from torch.optim.lr_scheduler import _LRScheduler
|
|
|
|
|
|
class MultiStepRestartLR(_LRScheduler):
|
|
""" MultiStep with restarts learning rate scheme.
|
|
|
|
Args:
|
|
optimizer (torch.nn.optimizer): Torch optimizer.
|
|
milestones (list): Iterations that will decrease learning rate.
|
|
gamma (float): Decrease ratio. Default: 0.1.
|
|
restarts (list): Restart iterations. Default: [0].
|
|
restart_weights (list): Restart weights at each restart iteration.
|
|
Default: [1].
|
|
last_epoch (int): Used in _LRScheduler. Default: -1.
|
|
"""
|
|
|
|
def __init__(self,
|
|
optimizer,
|
|
milestones,
|
|
gamma=0.1,
|
|
restarts=(0,),
|
|
restart_weights=(1,),
|
|
last_epoch=-1):
|
|
self.milestones = Counter(milestones)
|
|
self.gamma = gamma
|
|
self.restarts = restarts
|
|
self.restart_weights = restart_weights
|
|
print(type(self.restarts),self.restarts)
|
|
print(type(self.restart_weights),self.restart_weights)
|
|
assert len(self.restarts) == len(
|
|
self.restart_weights), 'restarts and their weights do not match.'
|
|
super(MultiStepRestartLR, self).__init__(optimizer, last_epoch)
|
|
|
|
def get_lr(self):
|
|
if self.last_epoch in self.restarts:
|
|
weight = self.restart_weights[self.restarts.index(self.last_epoch)]
|
|
return [
|
|
group['initial_lr'] * weight
|
|
for group in self.optimizer.param_groups
|
|
]
|
|
if self.last_epoch not in self.milestones:
|
|
return [group['lr'] for group in self.optimizer.param_groups]
|
|
return [
|
|
group['lr'] * self.gamma**self.milestones[self.last_epoch]
|
|
for group in self.optimizer.param_groups
|
|
]
|
|
|
|
|
|
def get_position_from_periods(iteration, cumulative_period):
|
|
"""Get the position from a period list.
|
|
|
|
It will return the index of the right-closest number in the period list.
|
|
For example, the cumulative_period = [100, 200, 300, 400],
|
|
if iteration == 50, return 0;
|
|
if iteration == 210, return 2;
|
|
if iteration == 300, return 2.
|
|
|
|
Args:
|
|
iteration (int): Current iteration.
|
|
cumulative_period (list[int]): Cumulative period list.
|
|
|
|
Returns:
|
|
int: The position of the right-closest number in the period list.
|
|
"""
|
|
for i, period in enumerate(cumulative_period):
|
|
if iteration <= period:
|
|
return i
|
|
|
|
|
|
class CosineAnnealingRestartLR(_LRScheduler):
|
|
""" Cosine annealing with restarts learning rate scheme.
|
|
|
|
An example of config:
|
|
periods = [10, 10, 10, 10]
|
|
restart_weights = [1, 0.5, 0.5, 0.5]
|
|
eta_min=1e-7
|
|
|
|
It has four cycles, each has 10 iterations. At 10th, 20th, 30th, the
|
|
scheduler will restart with the weights in restart_weights.
|
|
|
|
Args:
|
|
optimizer (torch.nn.optimizer): Torch optimizer.
|
|
periods (list): Period for each cosine anneling cycle.
|
|
restart_weights (list): Restart weights at each restart iteration.
|
|
Default: [1].
|
|
eta_min (float): The mimimum lr. Default: 0.
|
|
last_epoch (int): Used in _LRScheduler. Default: -1.
|
|
"""
|
|
|
|
def __init__(self,
|
|
optimizer,
|
|
periods,
|
|
restart_weights=(1),
|
|
eta_min=0,
|
|
last_epoch=-1):
|
|
self.periods = periods
|
|
self.restart_weights = restart_weights
|
|
self.eta_min = eta_min
|
|
assert (len(self.periods) == len(self.restart_weights)
|
|
), 'periods and restart_weights should have the same length.'
|
|
self.cumulative_period = [
|
|
sum(self.periods[0:i + 1]) for i in range(0, len(self.periods))
|
|
]
|
|
super(CosineAnnealingRestartLR, self).__init__(optimizer, last_epoch)
|
|
|
|
def get_lr(self):
|
|
idx = get_position_from_periods(self.last_epoch,
|
|
self.cumulative_period)
|
|
current_weight = self.restart_weights[idx]
|
|
nearest_restart = 0 if idx == 0 else self.cumulative_period[idx - 1]
|
|
current_period = self.periods[idx]
|
|
|
|
return [
|
|
self.eta_min + current_weight * 0.5 * (base_lr - self.eta_min) *
|
|
(1 + math.cos(math.pi * (
|
|
(self.last_epoch - nearest_restart) / current_period)))
|
|
for base_lr in self.base_lrs
|
|
] |