def test_get_optimizer(self): model = TempModel() for opt_name in AVAILABLE_OPTIMIZERS.keys(): opt_cls = optim.get_optimizer(opt_name) opt = opt_cls(model.parameters(), lr=self.INITIAL_LR) assert isinstance(opt, AVAILABLE_OPTIMIZERS[opt_name])
def test_get_optimizer(self): model = TempModel() for opt_name in AVAILABLE_OPTIMIZERS.keys(): if opt_name == 'fused_adam': if not torch.cuda.is_available(): continue opt_cls = optim.get_optimizer(opt_name) opt = opt_cls(model.parameters(), lr=self.INITIAL_LR) assert isinstance(opt, AVAILABLE_OPTIMIZERS[opt_name])
def test_sched_config_parse_simple(self): model = TempModel() opt_cls = optim.get_optimizer('novograd') opt = opt_cls(model.parameters(), lr=self.INITIAL_LR) basic_sched_config = {'name': 'CosineAnnealing', 'max_steps': 10} scheduler_setup = optim.lr_scheduler.prepare_lr_scheduler(opt, basic_sched_config) assert isinstance(scheduler_setup['scheduler'], optim.lr_scheduler.CosineAnnealing) dict_config = omegaconf.OmegaConf.create(basic_sched_config) scheduler_setup = optim.lr_scheduler.prepare_lr_scheduler(opt, dict_config) assert isinstance(scheduler_setup['scheduler'], optim.lr_scheduler.CosineAnnealing)
def test_register_optimizer(self): class TempOpt(torch.optim.SGD): pass class TempOptParams(config.optimizers.SGDParams): pass optim.register_optimizer('TempOpt', TempOpt, TempOptParams) model = TempModel() opt_cls = optim.get_optimizer('TempOpt') opt = opt_cls(model.parameters(), lr=self.INITIAL_LR) assert isinstance(opt, TempOpt)
def test_get_optimizer(self): model = TempModel() for opt_name in AVAILABLE_OPTIMIZERS.keys(): if opt_name == 'fused_adam': if not torch.cuda.is_available(): continue opt_cls = optim.get_optimizer(opt_name) if opt_name == 'adafactor': # Adafactor's default mode uses relative_step without any lr. opt = opt_cls(model.parameters()) else: opt = opt_cls(model.parameters(), lr=self.INITIAL_LR) assert isinstance(opt, AVAILABLE_OPTIMIZERS[opt_name])
def test_sched_config_parse_from_cls(self): model = TempModel() opt_cls = optim.get_optimizer('novograd') opt = opt_cls(model.parameters(), lr=self.INITIAL_LR) basic_sched_config = { 'target': 'nemo.core.config.CosineAnnealingParams', 'params': {'min_lr': 0.1}, 'max_steps': self.MAX_STEPS, } scheduler_setup = optim.lr_scheduler.prepare_lr_scheduler(opt, basic_sched_config) assert isinstance(scheduler_setup['scheduler'], optim.lr_scheduler.CosineAnnealing) dict_config = omegaconf.OmegaConf.create(basic_sched_config) scheduler_setup = optim.lr_scheduler.prepare_lr_scheduler(opt, dict_config) assert isinstance(scheduler_setup['scheduler'], optim.lr_scheduler.CosineAnnealing)
def test_register_scheduler(self): class TempSched(optim.lr_scheduler.CosineAnnealing): pass class TempSchedParams(config.schedulers.CosineAnnealingParams): pass optim.lr_scheduler.register_scheduler('TempSched', TempSched, TempSchedParams) model = TempModel() opt_cls = optim.get_optimizer('novograd') opt = opt_cls(model.parameters(), lr=self.INITIAL_LR) sched_cls = optim.lr_scheduler.get_scheduler('TempSched') sched = sched_cls(opt, max_steps=self.MAX_STEPS) assert isinstance(sched, TempSched)
def test_SquareAnnealing(self): model = TempModel() opt_cls = optim.get_optimizer('novograd') opt = opt_cls(model.parameters(), lr=self.INITIAL_LR) # No warmup case policy = optim.lr_scheduler.SquareAnnealing(opt, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR) initial_lr = policy.get_last_lr()[0] assert initial_lr == self.INITIAL_LR for i in range(self.MAX_STEPS): assert policy.get_last_lr()[0] <= self.INITIAL_LR opt.step() policy.step() policy.step() final_lr = policy.get_last_lr()[0] assert final_lr == self.MIN_LR # Warmup steps available policy = optim.lr_scheduler.SquareAnnealing(opt, warmup_steps=5, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR) initial_lr = policy.get_last_lr()[0] assert initial_lr < self.INITIAL_LR for i in range(self.MAX_STEPS): if i <= 5: assert policy.get_last_lr()[0] <= self.INITIAL_LR else: assert policy.get_last_lr()[0] < self.INITIAL_LR opt.step() policy.step() policy.step() final_lr = policy.get_last_lr()[0] assert final_lr == self.MIN_LR
def test_CosineAnnealing_with_noop_steps(self): model = TempModel() opt_cls = optim.get_optimizer('novograd') opt = opt_cls(model.parameters(), lr=self.INITIAL_LR) # No warmup case policy = optim.lr_scheduler.CosineAnnealing(opt, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR) initial_lr = policy.get_last_lr()[0] assert initial_lr == self.INITIAL_LR update_steps = 0 for i in range(self.MAX_STEPS): assert policy.get_last_lr()[0] <= self.INITIAL_LR opt.step() policy.step() # Perform a No-Op for scheduler every 2 steps if i % 2 == 0: policy.last_epoch -= 1 else: update_steps += 1 policy.step() update_steps += 1 assert update_steps < self.MAX_STEPS final_lr = policy.get_last_lr()[0] assert final_lr > self.MIN_LR # update step = true number of updates performed after some number of skipped steps true_end_lr = policy._get_lr(step=update_steps)[0] assert final_lr == true_end_lr
def test_PolynomialDecayAnnealing(self): model = TempModel() opt_cls = optim.get_optimizer('novograd') opt = opt_cls(model.parameters(), lr=self.INITIAL_LR) # No warmup case policy = optim.lr_scheduler.PolynomialHoldDecayAnnealing( opt, power=2, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR) initial_lr = policy.get_last_lr()[0] assert initial_lr == self.INITIAL_LR for i in range(self.MAX_STEPS): assert policy.get_last_lr()[0] <= self.INITIAL_LR opt.step() policy.step() policy.step() final_lr = policy.get_last_lr()[0] assert final_lr == self.MIN_LR # Warmup steps available policy = optim.lr_scheduler.PolynomialHoldDecayAnnealing( opt, power=2, warmup_steps=5, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR) initial_lr = policy.get_last_lr()[0] assert initial_lr < self.INITIAL_LR for i in range(self.MAX_STEPS): if i <= 5: assert policy.get_last_lr()[0] <= self.INITIAL_LR else: assert policy.get_last_lr()[0] < self.INITIAL_LR opt.step() policy.step() policy.step() final_lr = policy.get_last_lr()[0] assert final_lr == self.MIN_LR # Warmup + Hold steps available policy = optim.lr_scheduler.PolynomialHoldDecayAnnealing( opt, warmup_steps=5, hold_steps=3, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR, power=2) initial_lr = policy.get_last_lr()[0] assert initial_lr < self.INITIAL_LR for i in range(self.MAX_STEPS): if i <= 4: assert policy.get_last_lr()[0] <= self.INITIAL_LR elif i <= 8: assert policy.get_last_lr()[0] == self.INITIAL_LR else: assert policy.get_last_lr()[0] < self.INITIAL_LR opt.step() policy.step() policy.step() final_lr = policy.get_last_lr()[0] assert final_lr == self.MIN_LR