##################################################### # Copyright (c) Xuanyi Dong [GitHub D-X-Y], 2019.01 # ##################################################### import math, torch import torch.nn as nn from bisect import bisect_right from torch.optim import Optimizer class _LRScheduler(object): def __init__(self, optimizer, warmup_epochs, epochs): if not isinstance(optimizer, Optimizer): raise TypeError('{:} is not an Optimizer'.format(type(optimizer).__name__)) self.optimizer = optimizer for group in optimizer.param_groups: group.setdefault('initial_lr', group['lr']) self.base_lrs = list(map(lambda group: group['initial_lr'], optimizer.param_groups)) self.max_epochs = epochs self.warmup_epochs = warmup_epochs self.current_epoch = 0 self.current_iter = 0 def extra_repr(self): return '' def __repr__(self): return ('{name}(warmup={warmup_epochs}, max-epoch={max_epochs}, current::epoch={current_epoch}, iter={current_iter:.2f}'.format(name=self.__class__.__name__, **self.__dict__) + ', {:})'.format(self.extra_repr())) def state_dict(self): return {key: value for key, value in self.__dict__.items() if key != 'optimizer'} def load_state_dict(self, state_dict): self.__dict__.update(state_dict) def get_lr(self): raise NotImplementedError def get_min_info(self): lrs = self.get_lr() return '#LR=[{:.6f}~{:.6f}] epoch={:03d}, iter={:4.2f}#'.format(min(lrs), max(lrs), self.current_epoch, self.current_iter) def get_min_lr(self): return min( self.get_lr() ) def update(self, cur_epoch, cur_iter): if cur_epoch is not None: assert isinstance(cur_epoch, int) and cur_epoch>=0, 'invalid cur-epoch : {:}'.format(cur_epoch) self.current_epoch = cur_epoch if cur_iter is not None: assert isinstance(cur_iter, float) and cur_iter>=0, 'invalid cur-iter : {:}'.format(cur_iter) self.current_iter = cur_iter for param_group, lr in zip(self.optimizer.param_groups, self.get_lr()): param_group['lr'] = lr class CosineAnnealingLR(_LRScheduler): def __init__(self, optimizer, warmup_epochs, epochs, T_max, eta_min): self.T_max = T_max self.eta_min = eta_min super(CosineAnnealingLR, self).__init__(optimizer, warmup_epochs, epochs) def extra_repr(self): return 'type={:}, T-max={:}, eta-min={:}'.format('cosine', self.T_max, self.eta_min) def get_lr(self): lrs = [] for base_lr in self.base_lrs: if self.current_epoch >= self.warmup_epochs and self.current_epoch < self.max_epochs: last_epoch = self.current_epoch - self.warmup_epochs #if last_epoch < self.T_max: #if last_epoch < self.max_epochs: lr = self.eta_min + (base_lr - self.eta_min) * (1 + math.cos(math.pi * last_epoch / self.T_max)) / 2 #else: # lr = self.eta_min + (base_lr - self.eta_min) * (1 + math.cos(math.pi * (self.T_max-1.0) / self.T_max)) / 2 elif self.current_epoch >= self.max_epochs: lr = self.eta_min else: lr = (self.current_epoch / self.warmup_epochs + self.current_iter / self.warmup_epochs) * base_lr lrs.append( lr ) return lrs class MultiStepLR(_LRScheduler): def __init__(self, optimizer, warmup_epochs, epochs, milestones, gammas): assert len(milestones) == len(gammas), 'invalid {:} vs {:}'.format(len(milestones), len(gammas)) self.milestones = milestones self.gammas = gammas super(MultiStepLR, self).__init__(optimizer, warmup_epochs, epochs) def extra_repr(self): return 'type={:}, milestones={:}, gammas={:}, base-lrs={:}'.format('multistep', self.milestones, self.gammas, self.base_lrs) def get_lr(self): lrs = [] for base_lr in self.base_lrs: if self.current_epoch >= self.warmup_epochs: last_epoch = self.current_epoch - self.warmup_epochs idx = bisect_right(self.milestones, last_epoch) lr = base_lr for x in self.gammas[:idx]: lr *= x else: lr = (self.current_epoch / self.warmup_epochs + self.current_iter / self.warmup_epochs) * base_lr lrs.append( lr ) return lrs class ExponentialLR(_LRScheduler): def __init__(self, optimizer, warmup_epochs, epochs, gamma): self.gamma = gamma super(ExponentialLR, self).__init__(optimizer, warmup_epochs, epochs) def extra_repr(self): return 'type={:}, gamma={:}, base-lrs={:}'.format('exponential', self.gamma, self.base_lrs) def get_lr(self): lrs = [] for base_lr in self.base_lrs: if self.current_epoch >= self.warmup_epochs: last_epoch = self.current_epoch - self.warmup_epochs assert last_epoch >= 0, 'invalid last_epoch : {:}'.format(last_epoch) lr = base_lr * (self.gamma ** last_epoch) else: lr = (self.current_epoch / self.warmup_epochs + self.current_iter / self.warmup_epochs) * base_lr lrs.append( lr ) return lrs class LinearLR(_LRScheduler): def __init__(self, optimizer, warmup_epochs, epochs, max_LR, min_LR): self.max_LR = max_LR self.min_LR = min_LR super(LinearLR, self).__init__(optimizer, warmup_epochs, epochs) def extra_repr(self): return 'type={:}, max_LR={:}, min_LR={:}, base-lrs={:}'.format('LinearLR', self.max_LR, self.min_LR, self.base_lrs) def get_lr(self): lrs = [] for base_lr in self.base_lrs: if self.current_epoch >= self.warmup_epochs: last_epoch = self.current_epoch - self.warmup_epochs assert last_epoch >= 0, 'invalid last_epoch : {:}'.format(last_epoch) ratio = (self.max_LR - self.min_LR) * last_epoch / self.max_epochs / self.max_LR lr = base_lr * (1-ratio) else: lr = (self.current_epoch / self.warmup_epochs + self.current_iter / self.warmup_epochs) * base_lr lrs.append( lr ) return lrs class CrossEntropyLabelSmooth(nn.Module): def __init__(self, num_classes, epsilon): super(CrossEntropyLabelSmooth, self).__init__() self.num_classes = num_classes self.epsilon = epsilon self.logsoftmax = nn.LogSoftmax(dim=1) def forward(self, inputs, targets): log_probs = self.logsoftmax(inputs) targets = torch.zeros_like(log_probs).scatter_(1, targets.unsqueeze(1), 1) targets = (1 - self.epsilon) * targets + self.epsilon / self.num_classes loss = (-targets * log_probs).mean(0).sum() return loss def get_optim_scheduler(parameters, config): assert hasattr(config, 'optim') and hasattr(config, 'scheduler') and hasattr(config, 'criterion'), 'config must have optim / scheduler / criterion keys instead of {:}'.format(config) if config.optim == 'SGD': optim = torch.optim.SGD(parameters, config.LR, momentum=config.momentum, weight_decay=config.decay, nesterov=config.nesterov) elif config.optim == 'RMSprop': optim = torch.optim.RMSprop(parameters, config.LR, momentum=config.momentum, weight_decay=config.decay) else: raise ValueError('invalid optim : {:}'.format(config.optim)) if config.scheduler == 'cos': T_max = getattr(config, 'T_max', config.epochs) scheduler = CosineAnnealingLR(optim, config.warmup, config.epochs, T_max, config.eta_min) elif config.scheduler == 'multistep': scheduler = MultiStepLR(optim, config.warmup, config.epochs, config.milestones, config.gammas) elif config.scheduler == 'exponential': scheduler = ExponentialLR(optim, config.warmup, config.epochs, config.gamma) elif config.scheduler == 'linear': scheduler = LinearLR(optim, config.warmup, config.epochs, config.LR, config.LR_min) else: raise ValueError('invalid scheduler : {:}'.format(config.scheduler)) if config.criterion == 'Softmax': criterion = torch.nn.CrossEntropyLoss() elif config.criterion == 'SmoothSoftmax': criterion = CrossEntropyLabelSmooth(config.class_num, config.label_smooth) else: raise ValueError('invalid criterion : {:}'.format(config.criterion)) return optim, scheduler, criterion