Spaces:
Running
on
T4
Running
on
T4
| # -------------------------------------------------------- | |
| # SiamMask | |
| # Licensed under The MIT License | |
| # Written by Qiang Wang (wangqiang2015 at ia.ac.cn) | |
| # -------------------------------------------------------- | |
| from __future__ import division | |
| import numpy as np | |
| import math | |
| from torch.optim.lr_scheduler import _LRScheduler | |
| class LRScheduler(_LRScheduler): | |
| def __init__(self, optimizer, last_epoch=-1): | |
| if 'lr_spaces' not in self.__dict__: | |
| raise Exception('lr_spaces must be set in "LRSchduler"') | |
| super(LRScheduler, self).__init__(optimizer, last_epoch) | |
| def get_cur_lr(self): | |
| return self.lr_spaces[self.last_epoch] | |
| def get_lr(self): | |
| epoch = self.last_epoch | |
| return [self.lr_spaces[epoch] * pg['initial_lr'] / self.start_lr for pg in self.optimizer.param_groups] | |
| def __repr__(self): | |
| return "({}) lr spaces: \n{}".format(self.__class__.__name__, self.lr_spaces) | |
| class LogScheduler(LRScheduler): | |
| def __init__(self, optimizer, start_lr=0.03, end_lr=5e-4, epochs=50, last_epoch=-1, **kwargs): | |
| self.start_lr = start_lr | |
| self.end_lr = end_lr | |
| self.epochs = epochs | |
| self.lr_spaces = np.logspace(math.log10(start_lr), math.log10(end_lr), epochs) | |
| super(LogScheduler, self).__init__(optimizer, last_epoch) | |
| class StepScheduler(LRScheduler): | |
| def __init__(self, optimizer, start_lr=0.01, end_lr=None, step=10, mult=0.1, epochs=50, last_epoch=-1, **kwargs): | |
| if end_lr is not None: | |
| if start_lr is None: | |
| start_lr = end_lr / (mult ** (epochs // step)) | |
| else: # for warm up policy | |
| mult = math.pow(end_lr/start_lr, 1. / (epochs // step)) | |
| self.start_lr = start_lr | |
| self.lr_spaces = self.start_lr * (mult**(np.arange(epochs) // step)) | |
| self.mult = mult | |
| self._step = step | |
| super(StepScheduler, self).__init__(optimizer, last_epoch) | |
| class MultiStepScheduler(LRScheduler): | |
| def __init__(self, optimizer, start_lr=0.01, end_lr=None, steps=[10,20,30,40], mult=0.5, epochs=50, last_epoch=-1, **kwargs): | |
| if end_lr is not None: | |
| if start_lr is None: | |
| start_lr = end_lr / (mult ** (len(steps))) | |
| else: | |
| mult = math.pow(end_lr/start_lr, 1. / len(steps)) | |
| self.start_lr = start_lr | |
| self.lr_spaces = self._build_lr(start_lr, steps, mult, epochs) | |
| self.mult = mult | |
| self.steps = steps | |
| super(MultiStepScheduler, self).__init__(optimizer, last_epoch) | |
| def _build_lr(self, start_lr, steps, mult, epochs): | |
| lr = [0] * epochs | |
| lr[0] = start_lr | |
| for i in range(1, epochs): | |
| lr[i] = lr[i-1] | |
| if i in steps: | |
| lr[i] *= mult | |
| return np.array(lr, dtype=np.float32) | |
| class LinearStepScheduler(LRScheduler): | |
| def __init__(self, optimizer, start_lr=0.01, end_lr=0.005, epochs=50, last_epoch=-1, **kwargs): | |
| self.start_lr = start_lr | |
| self.end_lr = end_lr | |
| self.lr_spaces = np.linspace(start_lr, end_lr, epochs) | |
| super(LinearStepScheduler, self).__init__(optimizer, last_epoch) | |
| class CosStepScheduler(LRScheduler): | |
| def __init__(self, optimizer, start_lr=0.01, end_lr=0.005, epochs=50, last_epoch=-1, **kwargs): | |
| self.start_lr = start_lr | |
| self.end_lr = end_lr | |
| self.lr_spaces = self._build_lr(start_lr, end_lr, epochs) | |
| super(CosStepScheduler, self).__init__(optimizer, last_epoch) | |
| def _build_lr(self, start_lr, end_lr, epochs): | |
| index = np.arange(epochs).astype(np.float32) | |
| lr = end_lr + (start_lr - end_lr) * (1. + np.cos(index * np.pi/ epochs)) * 0.5 | |
| return lr.astype(np.float32) | |
| class WarmUPScheduler(LRScheduler): | |
| def __init__(self, optimizer, warmup, normal, epochs=50, last_epoch=-1): | |
| warmup = warmup.lr_spaces # [::-1] | |
| normal = normal.lr_spaces | |
| self.lr_spaces = np.concatenate([warmup, normal]) | |
| self.start_lr = normal[0] | |
| super(WarmUPScheduler, self).__init__(optimizer, last_epoch) | |
| LRs = { | |
| 'log': LogScheduler, | |
| 'step': StepScheduler, | |
| 'multi-step': MultiStepScheduler, | |
| 'linear': LinearStepScheduler, | |
| 'cos': CosStepScheduler} | |
| def _build_lr_scheduler(optimizer, cfg, epochs=50, last_epoch=-1): | |
| if 'type' not in cfg: | |
| # return LogScheduler(optimizer, last_epoch=last_epoch, epochs=epochs) | |
| cfg['type'] = 'log' | |
| if cfg['type'] not in LRs: | |
| raise Exception('Unknown type of LR Scheduler "%s"'%cfg['type']) | |
| return LRs[cfg['type']](optimizer, last_epoch=last_epoch, epochs=epochs, **cfg) | |
| def _build_warm_up_scheduler(optimizer, cfg, epochs=50, last_epoch=-1): | |
| warmup_epoch = cfg['warmup']['epoch'] | |
| sc1 = _build_lr_scheduler(optimizer, cfg['warmup'], warmup_epoch, last_epoch) | |
| sc2 = _build_lr_scheduler(optimizer, cfg, epochs - warmup_epoch, last_epoch) | |
| return WarmUPScheduler(optimizer, sc1, sc2, epochs, last_epoch) | |
| def build_lr_scheduler(optimizer, cfg, epochs=50, last_epoch=-1): | |
| if 'warmup' in cfg: | |
| return _build_warm_up_scheduler(optimizer, cfg, epochs, last_epoch) | |
| else: | |
| return _build_lr_scheduler(optimizer, cfg, epochs, last_epoch) | |
| if __name__ == '__main__': | |
| import torch.nn as nn | |
| from torch.optim import SGD | |
| class Net(nn.Module): | |
| def __init__(self): | |
| super(Net, self).__init__() | |
| self.conv = nn.Conv2d(10, 10, kernel_size=3) | |
| net = Net().parameters() | |
| optimizer = SGD(net, lr=0.01) | |
| # test1 | |
| step = { | |
| 'type': 'step', | |
| 'start_lr': 0.01, | |
| 'step': 10, | |
| 'mult': 0.1 | |
| } | |
| lr = build_lr_scheduler(optimizer, step) | |
| print(lr) | |
| log = { | |
| 'type': 'log', | |
| 'start_lr': 0.03, | |
| 'end_lr': 5e-4, | |
| } | |
| lr = build_lr_scheduler(optimizer, log) | |
| print(lr) | |
| log = { | |
| 'type': 'multi-step', | |
| "start_lr": 0.01, | |
| "mult": 0.1, | |
| "steps": [10, 15, 20] | |
| } | |
| lr = build_lr_scheduler(optimizer, log) | |
| print(lr) | |
| cos = { | |
| "type": 'cos', | |
| 'start_lr': 0.01, | |
| 'end_lr': 0.0005, | |
| } | |
| lr = build_lr_scheduler(optimizer, cos) | |
| print(lr) | |
| step = { | |
| 'type': 'step', | |
| 'start_lr': 0.001, | |
| 'end_lr': 0.03, | |
| 'step': 1, | |
| } | |
| warmup = log.copy() | |
| warmup['warmup'] = step | |
| warmup['warmup']['epoch'] = 5 | |
| lr = build_lr_scheduler(optimizer, warmup, epochs=55) | |
| print(lr) | |
| lr.step() | |
| print(lr.last_epoch) | |
| lr.step(5) | |
| print(lr.last_epoch) | |