def setup_schedulers(self): """Set up schedulers.""" train_opt = self.opt['train'] scheduler_type = train_opt['scheduler'].pop('type') if scheduler_type in ['MultiStepLR', 'MultiStepRestartLR']: for optimizer in self.optimizers: self.schedulers.append( lr_scheduler.MultiStepRestartLR(optimizer, **train_opt['scheduler'])) elif scheduler_type == 'CosineAnnealingRestartLR': for optimizer in self.optimizers: self.schedulers.append( lr_scheduler.CosineAnnealingRestartLR( optimizer, **train_opt['scheduler'])) elif scheduler_type == 'TrueCosineAnnealingLR': print('..', 'cosineannealingLR') for optimizer in self.optimizers: self.schedulers.append( torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, **train_opt['scheduler'])) elif scheduler_type == 'LinearLR': for optimizer in self.optimizers: self.schedulers.append( lr_scheduler.LinearLR(optimizer, train_opt['total_iter'])) elif scheduler_type == 'VibrateLR': for optimizer in self.optimizers: self.schedulers.append( lr_scheduler.VibrateLR(optimizer, train_opt['total_iter'])) else: raise NotImplementedError( f'Scheduler {scheduler_type} is not implemented yet.')
def setup_schedulers(self): """Set up schedulers.""" train_opt = self.opt['train'] scheduler_type = train_opt['scheduler'].pop('type') if scheduler_type in ['MultiStepLR', 'MultiStepRestartLR']: for optimizer in self.optimizers: self.schedulers.append(lr_scheduler.MultiStepRestartLR(optimizer, **train_opt['scheduler'])) elif scheduler_type == 'CosineAnnealingRestartLR': for optimizer in self.optimizers: self.schedulers.append(lr_scheduler.CosineAnnealingRestartLR(optimizer, **train_opt['scheduler'])) else: raise NotImplementedError(f'Scheduler {scheduler_type} is not implemented yet.')