meng2003/music2dance

1from torch.optim import lr_scheduler2import torch3from .nero import Nero4# import torch_optimizer as optim5import ast6from madgrad import MADGRAD7from pl_bolts.optimizers.lr_scheduler import LinearWarmupCosineAnnealingLR8 9 10def get_optimizers(net, opt):11    if opt.optimizer == "adam":12        optimizer = torch.optim.Adam(net.parameters(), lr=opt.learning_rate, weight_decay=opt.weight_decay)13    elif opt.optimizer == "adamw":14        optimizer = torch.optim.AdamW(net.parameters(), lr=opt.learning_rate, weight_decay=opt.weight_decay, eps=1e-05, betas=(0.9, 0.95))15    elif opt.optimizer == "sgd":16        optimizer = torch.optim.SGD(net.parameters(), lr=opt.learning_rate, momentum=opt.momentum, weight_decay=opt.weight_decay)17    elif opt.optimizer == "adagrad":18        optimizer = torch.optim.Adagrad(net.parameters(), lr=opt.learning_rate, weight_decay=opt.weight_decay)19    elif opt.optimizer == "adadelta":20        optimizer = torch.optim.Adadelta(net.parameters(), lr=opt.learning_rate, weight_decay=opt.weight_decay)21    elif opt.optimizer == "rmsprop":22        optimizer = torch.optim.Rmsprop(net.parameters(), lr=opt.learning_rate, weight_decay=opt.weight_decay)23    elif opt.optimizer == "nero":24        optimizer = Nero(net.parameters(), lr=opt.learning_rate)25    elif opt.optimizer == "madgrad":26        optimizer = MADGRAD(net.parameters(), lr=opt.learning_rate, weight_decay=opt.weight_decay, momentum=opt.momentum)27    # elif opt.optimizer == "ranger":28    #     optimizer = optim.Ranger(net.parameters(), lr=opt.learning_rate, alpha=0.5, k=6, N_sma_threshhold=5, betas=(.95, 0.999), eps=1e-5, weight_decay=0 )29    else:30        return NotImplementedError('optimizer [%s] is not implemented', opt.optimizer)31    return [optimizer]32 33def get_scheduler(optimizer, opt):34    if opt.lr_policy == 'lambda':35        def lambda_rule(epoch):36            nepochs = opt.max_epochs - opt.nepoch_decay #number of epochs before beginning to decay37            lr_l = 1.0 - max(0, epoch + opt.epoch_count - nepochs) / float(opt.nepoch_decay + 1)38            return lr_l39        scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)40    elif opt.lr_policy == 'exponential':41        scheduler = lr_scheduler.ExponentialLR(optimizer = optimizer, gamma = opt.lr_decay_factor)42    elif opt.lr_policy == 'step':43        scheduler = lr_scheduler.StepLR(optimizer, step_size=opt.lr_decay_iters, gamma=opt.lr_decay_factor)44    elif opt.lr_policy == 'multistep':45        scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=ast.literal_eval(opt.lr_decay_milestones), gamma=opt.lr_decay_factor)46    elif opt.lr_policy == 'plateau':47        scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, threshold=0.01, patience=5)48    elif opt.lr_policy == 'cosine':49        scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=opt.nepoch, eta_min=0)50    elif opt.lr_policy == 'cyclic':51        scheduler = CyclicLR(optimizer, base_lr=opt.learning_rate / 10, max_lr=opt.learning_rate,52                             step_size=opt.nepoch_decay, mode='triangular2')53    elif opt.lr_policy == 'reduceOnPlateau':54        scheduler = ReduceLROnPlateau(optimizer, 'min', factor=0.2)55    elif opt.lr_policy == 'LinearWarmupCosineAnnealing':56        scheduler = LinearWarmupCosineAnnealingLR(optimizer, warmup_epochs=opt.warmup_epochs, max_epochs=opt.max_epochs)57    else:58        return NotImplementedError('learning rate policy [%s] is not implemented', opt.lr_policy)59    return scheduler60 61 62class CyclicLR(object):63    """Sets the learning rate of each parameter group according to64    cyclical learning rate policy (CLR). The policy cycles the learning65    rate between two boundaries with a constant frequency, as detailed in66    the paper `Cyclical Learning Rates for Training Neural Networks`_.67    The distance between the two boundaries can be scaled on a per-iteration68    or per-cycle basis.69 70    Cyclical learning rate policy changes the learning rate after every batch.71    `batch_step` should be called after a batch has been used for training.72    To resume training, save `last_batch_iteration` and use it to instantiate `CycleLR`.73 74    This class has three built-in policies, as put forth in the paper:75    "triangular":76        A basic triangular cycle w/ no amplitude scaling.77    "triangular2":78        A basic triangular cycle that scales initial amplitude by half each cycle.79    "exp_range":80        A cycle that scales initial amplitude by gamma**(cycle iterations) at each81        cycle iteration.82 83    This implementation was adapted from the github repo: `bckenstler/CLR`_84 85    Args:86        optimizer (Optimizer): Wrapped optimizer.87        base_lr (float or list): Initial learning rate which is the88            lower boundary in the cycle for eachparam groups.89            Default: 0.00190        max_lr (float or list): Upper boundaries in the cycle for91            each parameter group. Functionally,92            it defines the cycle amplitude (max_lr - base_lr).93            The lr at any cycle is the sum of base_lr94            and some scaling of the amplitude; therefore95            max_lr may not actually be reached depending on96            scaling function. Default: 0.00697        step_size (int): Number of training iterations per98            half cycle. Authors suggest setting step_size99            2-8 x training iterations in epoch. Default: 2000100        mode (str): One of {triangular, triangular2, exp_range}.101            Values correspond to policies detailed above.102            If scale_fn is not None, this argument is ignored.103            Default: 'triangular'104        gamma (float): Constant in 'exp_range' scaling function:105            gamma**(cycle iterations)106            Default: 1.0107        scale_fn (function): Custom scaling policy defined by a single108            argument lambda function, where109            0 <= scale_fn(x) <= 1 for all x >= 0.110            mode paramater is ignored111            Default: None112        scale_mode (str): {'cycle', 'iterations'}.113            Defines whether scale_fn is evaluated on114            cycle number or cycle iterations (training115            iterations since start of cycle).116            Default: 'cycle'117        last_batch_iteration (int): The index of the last batch. Default: -1118 119    Example:120        >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)121        >>> scheduler = torch.optim.CyclicLR(optimizer)122        >>> data_loader = torch.utils.data.DataLoader(...)123        >>> for epoch in range(10):124        >>>     for batch in data_loader:125        >>>         scheduler.batch_step()126        >>>         train_batch(...)127 128    .. _Cyclical Learning Rates for Training Neural Networks: https://arxiv.org/abs/1506.01186129    .. _bckenstler/CLR: https://github.com/bckenstler/CLR130    """131 132    def __init__(self, optimizer, base_lr=1e-3, max_lr=6e-3,133                 step_size=2000, mode='triangular', gamma=1.,134                 scale_fn=None, scale_mode='cycle', last_batch_iteration=-1):135 136        if not isinstance(optimizer, Optimizer):137            raise TypeError('{} is not an Optimizer'.format(138                type(optimizer).__name__))139        self.optimizer = optimizer140 141        if isinstance(base_lr, list) or isinstance(base_lr, tuple):142            if len(base_lr) != len(optimizer.param_groups):143                raise ValueError("expected {} base_lr, got {}".format(144                    len(optimizer.param_groups), len(base_lr)))145            self.base_lrs = list(base_lr)146        else:147            self.base_lrs = [base_lr] * len(optimizer.param_groups)148 149        if isinstance(max_lr, list) or isinstance(max_lr, tuple):150            if len(max_lr) != len(optimizer.param_groups):151                raise ValueError("expected {} max_lr, got {}".format(152                    len(optimizer.param_groups), len(max_lr)))153            self.max_lrs = list(max_lr)154        else:155            self.max_lrs = [max_lr] * len(optimizer.param_groups)156 157        self.step_size = step_size158 159        if mode not in ['triangular', 'triangular2', 'exp_range'] \160                and scale_fn is None:161            raise ValueError('mode is invalid and scale_fn is None')162 163        self.mode = mode164        self.gamma = gamma165 166        if scale_fn is None:167            if self.mode == 'triangular':168                self.scale_fn = self._triangular_scale_fn169                self.scale_mode = 'cycle'170            elif self.mode == 'triangular2':171                self.scale_fn = self._triangular2_scale_fn172                self.scale_mode = 'cycle'173            elif self.mode == 'exp_range':174                self.scale_fn = self._exp_range_scale_fn175                self.scale_mode = 'iterations'176        else:177            self.scale_fn = scale_fn178            self.scale_mode = scale_mode179 180        self.batch_step(last_batch_iteration + 1)181        self.last_batch_iteration = last_batch_iteration182 183    def batch_step(self, batch_iteration=None):184        if batch_iteration is None:185            batch_iteration = self.last_batch_iteration + 1186        self.last_batch_iteration = batch_iteration187        for param_group, lr in zip(self.optimizer.param_groups, self.get_lr()):188            param_group['lr'] = lr189 190    def _triangular_scale_fn(self, x):191        return 1.192 193    def _triangular2_scale_fn(self, x):194        return 1 / (2. ** (x - 1))195 196    def _exp_range_scale_fn(self, x):197        return self.gamma**(x)198 199    def get_lr(self):200        step_size = float(self.step_size)201        cycle = np.floor(1 + self.last_batch_iteration / (2 * step_size))202        x = np.abs(self.last_batch_iteration / step_size - 2 * cycle + 1)203 204        lrs = []205        param_lrs = zip(self.optimizer.param_groups, self.base_lrs, self.max_lrs)206        for param_group, base_lr, max_lr in param_lrs:207            base_height = (max_lr - base_lr) * np.maximum(0, (1 - x))208            if self.scale_mode == 'cycle':209                lr = base_lr + base_height * self.scale_fn(cycle)210            else:211                lr = base_lr + base_height * self.scale_fn(self.last_batch_iteration)212            lrs.append(lr)213        return lrs214