def test_tin_lr_updater_hook(): sys.modules['pavi'] = MagicMock() loader = DataLoader(torch.ones((10, 2))) runner = _build_demo_runner() hook_cfg = dict(type='TINLrUpdaterHook', min_lr=0.1) runner.register_hook_from_cfg(hook_cfg) hook_cfg = dict(type='TINLrUpdaterHook', by_epoch=False, min_lr=0.1, warmup='exp', warmup_iters=2, warmup_ratio=0.9) runner.register_hook_from_cfg(hook_cfg) runner.register_hook_from_cfg(dict(type='IterTimerHook')) runner.register_hook(IterTimerHook()) hook_cfg = dict(type='TINLrUpdaterHook', by_epoch=False, min_lr=0.1, warmup='constant', warmup_iters=2, warmup_ratio=0.9) runner.register_hook_from_cfg(hook_cfg) runner.register_hook_from_cfg(dict(type='IterTimerHook')) runner.register_hook(IterTimerHook()) hook_cfg = dict(type='TINLrUpdaterHook', by_epoch=False, min_lr=0.1, warmup='linear', warmup_iters=2, warmup_ratio=0.9) runner.register_hook_from_cfg(hook_cfg) runner.register_hook_from_cfg(dict(type='IterTimerHook')) runner.register_hook(IterTimerHook()) # add pavi hook hook = PaviLoggerHook(interval=1, add_graph=False, add_last_ckpt=True) runner.register_hook(hook) runner.run([loader], [('train', 1)]) shutil.rmtree(runner.work_dir) assert hasattr(hook, 'writer') calls = [ call('train', { 'learning_rate': 0.028544155877284292, 'momentum': 0.95 }, 1), call('train', { 'learning_rate': 0.04469266270539641, 'momentum': 0.95 }, 6), call('train', { 'learning_rate': 0.09695518130045147, 'momentum': 0.95 }, 10) ] hook.writer.add_scalars.assert_has_calls(calls, any_order=True)
def test_cosine_runner_hook(): """xdoctest -m tests/test_hooks.py test_cosine_runner_hook.""" sys.modules['pavi'] = MagicMock() loader = DataLoader(torch.ones((10, 2))) runner = _build_demo_runner() # add momentum scheduler hook_cfg = dict( type='CosineAnnealingMomentumUpdaterHook', min_momentum_ratio=0.99 / 0.95, by_epoch=False, warmup_iters=2, warmup_ratio=0.9 / 0.95) runner.register_hook_from_cfg(hook_cfg) # add momentum LR scheduler hook_cfg = dict( type='CosineAnnealingLrUpdaterHook', by_epoch=False, min_lr_ratio=0, warmup_iters=2, warmup_ratio=0.9) runner.register_hook_from_cfg(hook_cfg) runner.register_hook_from_cfg(dict(type='IterTimerHook')) runner.register_hook(IterTimerHook()) # add pavi hook hook = PaviLoggerHook(interval=1, add_graph=False, add_last_ckpt=True) runner.register_hook(hook) runner.run([loader], [('train', 1)]) shutil.rmtree(runner.work_dir) # TODO: use a more elegant way to check values assert hasattr(hook, 'writer') calls = [ call('train', { 'learning_rate': 0.02, 'momentum': 0.95 }, 0), call('train', { 'learning_rate': 0.01, 'momentum': 0.97 }, 5), call('train', { 'learning_rate': 0.0004894348370484647, 'momentum': 0.9890211303259032 }, 9) ] hook.writer.add_scalars.assert_has_calls(calls, any_order=True)
def test_momentum_runner_hook(): """ xdoctest -m tests/test_hooks.py test_momentum_runner_hook """ sys.modules['pavi'] = MagicMock() loader = DataLoader(torch.ones((10, 2))) runner = _build_demo_runner() # add momentum scheduler hook = CyclicMomentumUpdaterHook( by_epoch=False, target_ratio=(0.85 / 0.95, 1), cyclic_times=1, step_ratio_up=0.4) runner.register_hook(hook) # add momentum LR scheduler hook = CyclicLrUpdaterHook( by_epoch=False, target_ratio=(10, 1), cyclic_times=1, step_ratio_up=0.4) runner.register_hook(hook) runner.register_hook(IterTimerHook()) # add pavi hook hook = PaviLoggerHook(interval=1, add_graph=False, add_last_ckpt=True) runner.register_hook(hook) runner.run([loader], [('train', 1)], 1) shutil.rmtree(runner.work_dir) # TODO: use a more elegant way to check values assert hasattr(hook, 'writer') calls = [ call('train', { 'learning_rate': 0.01999999999999999, 'momentum': 0.95 }, 0), call('train', { 'learning_rate': 0.2, 'momentum': 0.85 }, 4), call('train', { 'learning_rate': 0.155, 'momentum': 0.875 }, 6), ] hook.writer.add_scalars.assert_has_calls(calls, any_order=True)
def register_training_hooks(self, lr_config, optimizer_config=None, checkpoint_config=None, log_config=None, momentum_config=None): """Register default hooks for training. Default hooks include: - LrUpdaterHook - MomentumUpdaterHook - OptimizerStepperHook - CheckpointSaverHook - IterTimerHook - LoggerHook(s) """ self.register_lr_hook(lr_config) self.register_momentum_hook(momentum_config) self.register_optimizer_hook(optimizer_config) self.register_checkpoint_hook(checkpoint_config) self.register_hook(IterTimerHook()) self.register_logger_hooks(log_config)
def register_qat_hooks(self, loss, metrics, lr_policies, qat_policies, ckpt_interval=None, runtime_hook=None): assert isinstance(loss, dict) assert isinstance(metrics, (tuple, list)) assert isinstance(lr_policies, (tuple, list)) assert isinstance(qat_policies, (tuple, list)) loss = training.build_loss(loss) metrics = training.build_metrics(*metrics) lr_policies = training.build_lr_policies(*lr_policies) qat_policies = training.build_qat_policies(*qat_policies) # make sure loss firstly getting ready after `batch_processor` self.register_hook(loss, priority="HIGH") self.register_hook(IterTimerHook()) if ckpt_interval: self.register_hook(CheckpointHook(interval=ckpt_interval)) for hook in chain(metrics, qat_policies, lr_policies): if isinstance(hook, HijackModuleOutput): priority = "LOW" else: priority = "NORMAL" self.register_hook(hook, priority) if runtime_hook is not None: interval = runtime_hook["interval"] hooks = runtime_hook["hooks"] post_process = runtime_hook.get("post_process") self.inject_runtime_hooks(interval, hooks, post_process) else: self.inject_runtime_hooks(-1, [], None)
def test_cyclic_lr_update_hook(multi_optimizers, max_iters): """Test CyclicLrUpdateHook.""" with pytest.raises(AssertionError): # by_epoch should be False CyclicLrUpdaterHook(by_epoch=True) with pytest.raises(AssertionError): # target_ratio" must be either float or tuple/list of two floats CyclicLrUpdaterHook(by_epoch=False, target_ratio=(10.0, 0.1, 0.2)) with pytest.raises(AssertionError): # step_ratio_up" must be in range [0,1) CyclicLrUpdaterHook(by_epoch=False, step_ratio_up=1.4) with pytest.raises(ValueError): # anneal_strategy must be one of "cos" or "linear" CyclicLrUpdaterHook(by_epoch=False, anneal_strategy='sin') sys.modules['pavi'] = MagicMock() loader = DataLoader(torch.ones((10, 2))) runner = _build_demo_runner(runner_type='IterBasedRunner', max_epochs=None, max_iters=max_iters, multi_optimziers=multi_optimizers) # add cyclic LR scheduler hook = CyclicLrUpdaterHook(by_epoch=False, target_ratio=(10.0, 1.0), cyclic_times=1, step_ratio_up=0.5, anneal_strategy='linear') runner.register_hook(hook) runner.register_hook_from_cfg(dict(type='IterTimerHook')) runner.register_hook(IterTimerHook()) # add pavi hook hook = PaviLoggerHook(interval=1, add_graph=False, add_last_ckpt=True) runner.register_hook(hook) runner.run([loader], [('train', 1)]) shutil.rmtree(runner.work_dir) assert hasattr(hook, 'writer') if multi_optimizers: calls = [ call( 'train', { 'learning_rate/model1': 0.02, 'learning_rate/model2': 0.01, 'momentum/model1': 0.95, 'momentum/model2': 0.9, }, 1), call( 'train', { 'learning_rate/model1': 0.155, 'learning_rate/model2': 0.0775, 'momentum/model1': 0.95, 'momentum/model2': 0.9, }, 4), call( 'train', { 'learning_rate/model1': 0.155, 'learning_rate/model2': 0.0775, 'momentum/model1': 0.95, 'momentum/model2': 0.9, }, 6) ] else: calls = [ call('train', { 'learning_rate': 0.02, 'momentum': 0.95 }, 1), call('train', { 'learning_rate': 0.155, 'momentum': 0.95 }, 4), call('train', { 'learning_rate': 0.155, 'momentum': 0.95 }, 6), ] hook.writer.add_scalars.assert_has_calls(calls, any_order=True)
def test_step_runner_hook(multi_optimziers): """Test StepLrUpdaterHook.""" with pytest.raises(TypeError): # `step` should be specified StepLrUpdaterHook() with pytest.raises(AssertionError): # if `step` is int, should be positive StepLrUpdaterHook(-10) with pytest.raises(AssertionError): # if `step` is list of int, should all be positive StepLrUpdaterHook([10, 16, -20]) # test StepLrUpdaterHook with int `step` value sys.modules['pavi'] = MagicMock() loader = DataLoader(torch.ones((30, 2))) runner = _build_demo_runner(multi_optimziers=multi_optimziers) # add momentum scheduler hook_cfg = dict(type='StepMomentumUpdaterHook', by_epoch=False, step=5, gamma=0.5, min_momentum=0.05) runner.register_hook_from_cfg(hook_cfg) # add step LR scheduler hook = StepLrUpdaterHook(by_epoch=False, step=5, gamma=0.5, min_lr=1e-3) runner.register_hook(hook) runner.register_hook(IterTimerHook()) # add pavi hook hook = PaviLoggerHook(interval=1, add_graph=False, add_last_ckpt=True) runner.register_hook(hook) runner.run([loader], [('train', 1)]) shutil.rmtree(runner.work_dir) # TODO: use a more elegant way to check values assert hasattr(hook, 'writer') if multi_optimziers: calls = [ call( 'train', { 'learning_rate/model1': 0.02, 'learning_rate/model2': 0.01, 'momentum/model1': 0.95, 'momentum/model2': 0.9 }, 1), call( 'train', { 'learning_rate/model1': 0.01, 'learning_rate/model2': 0.005, 'momentum/model1': 0.475, 'momentum/model2': 0.45 }, 6), call( 'train', { 'learning_rate/model1': 0.0025, 'learning_rate/model2': 0.00125, 'momentum/model1': 0.11875, 'momentum/model2': 0.1125 }, 16), call( 'train', { 'learning_rate/model1': 0.00125, 'learning_rate/model2': 0.001, 'momentum/model1': 0.059375, 'momentum/model2': 0.05625 }, 21), call( 'train', { 'learning_rate/model1': 0.001, 'learning_rate/model2': 0.001, 'momentum/model1': 0.05, 'momentum/model2': 0.05 }, 26), call( 'train', { 'learning_rate/model1': 0.001, 'learning_rate/model2': 0.001, 'momentum/model1': 0.05, 'momentum/model2': 0.05 }, 30) ] else: calls = [ call('train', { 'learning_rate': 0.02, 'momentum': 0.95 }, 1), call('train', { 'learning_rate': 0.01, 'momentum': 0.475 }, 6), call('train', { 'learning_rate': 0.0025, 'momentum': 0.11875 }, 16), call('train', { 'learning_rate': 0.00125, 'momentum': 0.059375 }, 21), call('train', { 'learning_rate': 0.001, 'momentum': 0.05 }, 26), call('train', { 'learning_rate': 0.001, 'momentum': 0.05 }, 30) ] hook.writer.add_scalars.assert_has_calls(calls, any_order=True) # test StepLrUpdaterHook with list[int] `step` value sys.modules['pavi'] = MagicMock() loader = DataLoader(torch.ones((10, 2))) runner = _build_demo_runner(multi_optimziers=multi_optimziers) # add momentum scheduler hook_cfg = dict(type='StepMomentumUpdaterHook', by_epoch=False, step=[4, 6, 8], gamma=0.1) runner.register_hook_from_cfg(hook_cfg) # add step LR scheduler hook = StepLrUpdaterHook(by_epoch=False, step=[4, 6, 8], gamma=0.1) runner.register_hook(hook) runner.register_hook(IterTimerHook()) # add pavi hook hook = PaviLoggerHook(interval=1, add_graph=False, add_last_ckpt=True) runner.register_hook(hook) runner.run([loader], [('train', 1)]) shutil.rmtree(runner.work_dir) # TODO: use a more elegant way to check values assert hasattr(hook, 'writer') if multi_optimziers: calls = [ call( 'train', { 'learning_rate/model1': 0.02, 'learning_rate/model2': 0.01, 'momentum/model1': 0.95, 'momentum/model2': 0.9 }, 1), call( 'train', { 'learning_rate/model1': 0.002, 'learning_rate/model2': 0.001, 'momentum/model1': 9.5e-2, 'momentum/model2': 9.000000000000001e-2 }, 5), call( 'train', { 'learning_rate/model1': 2.0000000000000004e-4, 'learning_rate/model2': 1.0000000000000002e-4, 'momentum/model1': 9.500000000000001e-3, 'momentum/model2': 9.000000000000003e-3 }, 7), call( 'train', { 'learning_rate/model1': 2.0000000000000005e-05, 'learning_rate/model2': 1.0000000000000003e-05, 'momentum/model1': 9.500000000000002e-4, 'momentum/model2': 9.000000000000002e-4 }, 9) ] else: calls = [ call('train', { 'learning_rate': 0.02, 'momentum': 0.95 }, 1), call('train', { 'learning_rate': 0.002, 'momentum': 0.095 }, 5), call( 'train', { 'learning_rate': 2.0000000000000004e-4, 'momentum': 9.500000000000001e-3 }, 7), call( 'train', { 'learning_rate': 2.0000000000000005e-05, 'momentum': 9.500000000000002e-4 }, 9) ] hook.writer.add_scalars.assert_has_calls(calls, any_order=True)
def test_cosine_restart_lr_update_hook(multi_optimziers): """Test CosineRestartLrUpdaterHook.""" with pytest.raises(AssertionError): # either `min_lr` or `min_lr_ratio` should be specified CosineRestartLrUpdaterHook(by_epoch=False, periods=[2, 10], restart_weights=[0.5, 0.5], min_lr=0.1, min_lr_ratio=0) with pytest.raises(AssertionError): # periods and restart_weights should have the same length CosineRestartLrUpdaterHook(by_epoch=False, periods=[2, 10], restart_weights=[0.5], min_lr_ratio=0) with pytest.raises(ValueError): # the last cumulative_periods 7 (out of [5, 7]) should >= 10 sys.modules['pavi'] = MagicMock() loader = DataLoader(torch.ones((10, 2))) runner = _build_demo_runner() # add cosine restart LR scheduler hook = CosineRestartLrUpdaterHook( by_epoch=False, periods=[5, 2], # cumulative_periods [5, 7 (5 + 2)] restart_weights=[0.5, 0.5], min_lr=0.0001) runner.register_hook(hook) runner.register_hook(IterTimerHook()) # add pavi hook hook = PaviLoggerHook(interval=1, add_graph=False, add_last_ckpt=True) runner.register_hook(hook) runner.run([loader], [('train', 1)]) shutil.rmtree(runner.work_dir) sys.modules['pavi'] = MagicMock() loader = DataLoader(torch.ones((10, 2))) runner = _build_demo_runner(multi_optimziers=multi_optimziers) # add cosine restart LR scheduler hook = CosineRestartLrUpdaterHook(by_epoch=False, periods=[5, 5], restart_weights=[0.5, 0.5], min_lr_ratio=0) runner.register_hook(hook) runner.register_hook(IterTimerHook()) # add pavi hook hook = PaviLoggerHook(interval=1, add_graph=False, add_last_ckpt=True) runner.register_hook(hook) runner.run([loader], [('train', 1)]) shutil.rmtree(runner.work_dir) # TODO: use a more elegant way to check values assert hasattr(hook, 'writer') if multi_optimziers: calls = [ call( 'train', { 'learning_rate/model1': 0.01, 'learning_rate/model2': 0.005, 'momentum/model1': 0.95, 'momentum/model2': 0.9, }, 1), call( 'train', { 'learning_rate/model1': 0.01, 'learning_rate/model2': 0.005, 'momentum/model1': 0.95, 'momentum/model2': 0.9, }, 6), call( 'train', { 'learning_rate/model1': 0.0009549150281252633, 'learning_rate/model2': 0.00047745751406263163, 'momentum/model1': 0.95, 'momentum/model2': 0.9, }, 10) ] else: calls = [ call('train', { 'learning_rate': 0.01, 'momentum': 0.95 }, 1), call('train', { 'learning_rate': 0.01, 'momentum': 0.95 }, 6), call('train', { 'learning_rate': 0.0009549150281252633, 'momentum': 0.95 }, 10) ] hook.writer.add_scalars.assert_has_calls(calls, any_order=True)
def test_one_cycle_runner_hook(multi_optimziers, max_iters): """Test OneCycleLrUpdaterHook and OneCycleMomentumUpdaterHook.""" with pytest.raises(AssertionError): # by_epoch should be False OneCycleLrUpdaterHook(max_lr=0.1, by_epoch=True) with pytest.raises(ValueError): # expected float between 0 and 1 OneCycleLrUpdaterHook(max_lr=0.1, pct_start=-0.1) with pytest.raises(ValueError): # anneal_strategy should be either 'cos' or 'linear' OneCycleLrUpdaterHook(max_lr=0.1, anneal_strategy='sin') sys.modules['pavi'] = MagicMock() loader = DataLoader(torch.ones((10, 2))) runner = _build_demo_runner(multi_optimziers=multi_optimziers) # add momentum scheduler hook_cfg = dict(type='OneCycleMomentumUpdaterHook', base_momentum=0.85, max_momentum=0.95, pct_start=0.5, anneal_strategy='cos', three_phase=False) runner.register_hook_from_cfg(hook_cfg) # add LR scheduler hook_cfg = dict(type='OneCycleLrUpdaterHook', max_lr=0.01, pct_start=0.5, anneal_strategy='cos', div_factor=25, final_div_factor=1e4, three_phase=False) runner.register_hook_from_cfg(hook_cfg) runner.register_hook_from_cfg(dict(type='IterTimerHook')) runner.register_hook(IterTimerHook()) # add pavi hook hook = PaviLoggerHook(interval=1, add_graph=False, add_last_ckpt=True) runner.register_hook(hook) runner.run([loader], [('train', 1)]) shutil.rmtree(runner.work_dir) # TODO: use a more elegant way to check values assert hasattr(hook, 'writer') if multi_optimziers: calls = [ call( 'train', { 'learning_rate/model1': 0.0003999999999999993, 'learning_rate/model2': 0.0003999999999999993, 'momentum/model1': 0.95, 'momentum/model2': 0.95, }, 1), call( 'train', { 'learning_rate/model1': 0.00904508879153485, 'learning_rate/model2': 0.00904508879153485, 'momentum/model1': 0.8595491502812526, 'momentum/model2': 0.8595491502812526, }, 6), call( 'train', { 'learning_rate/model1': 4e-08, 'learning_rate/model2': 4e-08, 'momentum/model1': 0.95, 'momentum/model2': 0.95, }, 10) ] else: calls = [ call('train', { 'learning_rate': 0.0003999999999999993, 'momentum': 0.95 }, 1), call( 'train', { 'learning_rate': 0.00904508879153485, 'momentum': 0.8595491502812526 }, 6), call('train', { 'learning_rate': 4e-08, 'momentum': 0.95 }, 10) ] hook.writer.add_scalars.assert_has_calls(calls, any_order=True) # Test OneCycleLrUpdaterHook sys.modules['pavi'] = MagicMock() loader = DataLoader(torch.ones((10, 2))) runner = _build_demo_runner(runner_type='IterBasedRunner', max_epochs=None, max_iters=max_iters) args = dict( max_lr=0.01, total_steps=5, pct_start=0.5, anneal_strategy='linear', div_factor=25, final_div_factor=1e4, ) hook = OneCycleLrUpdaterHook(**args) runner.register_hook(hook) if max_iters == 10: # test total_steps < max_iters with pytest.raises(ValueError): runner.run([loader], [('train', 1)]) else: # test total_steps > max_iters runner.run([loader], [('train', 1)]) lr_last = runner.current_lr() t = torch.tensor([0.0], requires_grad=True) optim = torch.optim.SGD([t], lr=0.01) lr_scheduler = torch.optim.lr_scheduler.OneCycleLR(optim, **args) lr_target = [] for _ in range(max_iters): optim.step() lr_target.append(optim.param_groups[0]['lr']) lr_scheduler.step() assert lr_target[-1] == lr_last[0]
def test_yolox_lrupdater_hook(multi_optimziers): """xdoctest -m tests/test_hooks.py test_cosine_runner_hook.""" # Only used to prevent program errors YOLOXLrUpdaterHook(0, min_lr_ratio=0.05) sys.modules['pavi'] = MagicMock() loader = DataLoader(torch.ones((10, 2))) runner = _build_demo_runner(multi_optimziers=multi_optimziers) hook_cfg = dict( type='YOLOXLrUpdaterHook', warmup='exp', by_epoch=False, warmup_by_epoch=True, warmup_ratio=1, warmup_iters=5, # 5 epoch num_last_epochs=15, min_lr_ratio=0.05) runner.register_hook_from_cfg(hook_cfg) runner.register_hook_from_cfg(dict(type='IterTimerHook')) runner.register_hook(IterTimerHook()) # add pavi hook hook = PaviLoggerHook(interval=1, add_graph=False, add_last_ckpt=True) runner.register_hook(hook) runner.run([loader], [('train', 1)]) shutil.rmtree(runner.work_dir) # TODO: use a more elegant way to check values assert hasattr(hook, 'writer') if multi_optimziers: calls = [ call( 'train', { 'learning_rate/model1': 8.000000000000001e-06, 'learning_rate/model2': 4.000000000000001e-06, 'momentum/model1': 0.95, 'momentum/model2': 0.9 }, 1), call( 'train', { 'learning_rate/model1': 0.00039200000000000004, 'learning_rate/model2': 0.00019600000000000002, 'momentum/model1': 0.95, 'momentum/model2': 0.9 }, 7), call( 'train', { 'learning_rate/model1': 0.0008000000000000001, 'learning_rate/model2': 0.0004000000000000001, 'momentum/model1': 0.95, 'momentum/model2': 0.9 }, 10) ] else: calls = [ call('train', { 'learning_rate': 8.000000000000001e-06, 'momentum': 0.95 }, 1), call('train', { 'learning_rate': 0.00039200000000000004, 'momentum': 0.95 }, 7), call('train', { 'learning_rate': 0.0008000000000000001, 'momentum': 0.95 }, 10) ] hook.writer.add_scalars.assert_has_calls(calls, any_order=True)
def test_one_cycle_runner_hook(multi_optimziers): """Test OneCycleLrUpdaterHook and OneCycleMomentumUpdaterHook.""" with pytest.raises(AssertionError): # by_epoch should be False OneCycleLrUpdaterHook(max_lr=0.1, by_epoch=True) with pytest.raises(ValueError): # expected float between 0 and 1 OneCycleLrUpdaterHook(max_lr=0.1, pct_start=-0.1) with pytest.raises(ValueError): # anneal_strategy should be either 'cos' or 'linear' OneCycleLrUpdaterHook(max_lr=0.1, anneal_strategy='sin') sys.modules['pavi'] = MagicMock() loader = DataLoader(torch.ones((10, 2))) runner = _build_demo_runner(multi_optimziers=multi_optimziers) # add momentum scheduler hook_cfg = dict(type='OneCycleMomentumUpdaterHook', base_momentum=0.85, max_momentum=0.95, pct_start=0.5, anneal_strategy='cos', three_phase=False) runner.register_hook_from_cfg(hook_cfg) # add LR scheduler hook_cfg = dict(type='OneCycleLrUpdaterHook', max_lr=0.01, pct_start=0.5, anneal_strategy='cos', div_factor=25, final_div_factor=1e4, three_phase=False) runner.register_hook_from_cfg(hook_cfg) runner.register_hook_from_cfg(dict(type='IterTimerHook')) runner.register_hook(IterTimerHook()) # add pavi hook hook = PaviLoggerHook(interval=1, add_graph=False, add_last_ckpt=True) runner.register_hook(hook) runner.run([loader], [('train', 1)]) shutil.rmtree(runner.work_dir) # TODO: use a more elegant way to check values assert hasattr(hook, 'writer') if multi_optimziers: calls = [ call( 'train', { 'learning_rate/model1': 0.0003999999999999993, 'learning_rate/model2': 0.0003999999999999993, 'momentum/model1': 0.95, 'momentum/model2': 0.95, }, 1), call( 'train', { 'learning_rate/model1': 0.00904508879153485, 'learning_rate/model2': 0.00904508879153485, 'momentum/model1': 0.8595491502812526, 'momentum/model2': 0.8595491502812526, }, 6), call( 'train', { 'learning_rate/model1': 4e-08, 'learning_rate/model2': 4e-08, 'momentum/model1': 0.95, 'momentum/model2': 0.95, }, 10) ] else: calls = [ call('train', { 'learning_rate': 0.0003999999999999993, 'momentum': 0.95 }, 1), call( 'train', { 'learning_rate': 0.00904508879153485, 'momentum': 0.8595491502812526 }, 6), call('train', { 'learning_rate': 4e-08, 'momentum': 0.95 }, 10) ] hook.writer.add_scalars.assert_has_calls(calls, any_order=True)
def test_flat_cosine_runner_hook(multi_optimziers, by_epoch): """xdoctest -m tests/test_hooks.py test_flat_cosine_runner_hook.""" sys.modules['pavi'] = MagicMock() loader = DataLoader(torch.ones((10, 2))) max_epochs = 10 if by_epoch else 1 runner = _build_demo_runner(multi_optimziers=multi_optimziers, max_epochs=max_epochs) with pytest.raises(ValueError): # start_percent: expected float between 0 and 1 FlatCosineAnnealingLrUpdaterHook(start_percent=-0.1, min_lr_ratio=0) # add LR scheduler hook_cfg = dict(type='FlatCosineAnnealingLrUpdaterHook', by_epoch=by_epoch, min_lr_ratio=0, warmup='linear', warmup_iters=10 if by_epoch else 2, warmup_ratio=0.9, start_percent=0.5) runner.register_hook_from_cfg(hook_cfg) runner.register_hook_from_cfg(dict(type='IterTimerHook')) runner.register_hook(IterTimerHook()) # add pavi hook hook = PaviLoggerHook(interval=1, add_graph=False, add_last_ckpt=True) runner.register_hook(hook) runner.run([loader], [('train', 1)]) shutil.rmtree(runner.work_dir) # TODO: use a more elegant way to check values assert hasattr(hook, 'writer') if multi_optimziers: if by_epoch: calls = [ call( 'train', { 'learning_rate/model1': 0.018000000000000002, 'learning_rate/model2': 0.009000000000000001, 'momentum/model1': 0.95, 'momentum/model2': 0.9, }, 1), call( 'train', { 'learning_rate/model1': 0.02, 'learning_rate/model2': 0.01, 'momentum/model1': 0.95, 'momentum/model2': 0.9, }, 11), call( 'train', { 'learning_rate/model1': 0.018090169943749474, 'learning_rate/model2': 0.009045084971874737, 'momentum/model1': 0.95, 'momentum/model2': 0.9, }, 61), call( 'train', { 'learning_rate/model1': 0.0019098300562505265, 'learning_rate/model2': 0.0009549150281252633, 'momentum/model1': 0.95, 'momentum/model2': 0.9, }, 100) ] else: calls = [ call( 'train', { 'learning_rate/model1': 0.018000000000000002, 'learning_rate/model2': 0.009000000000000001, 'momentum/model1': 0.95, 'momentum/model2': 0.9 }, 1), call( 'train', { 'learning_rate/model1': 0.02, 'learning_rate/model2': 0.01, 'momentum/model1': 0.95, 'momentum/model2': 0.9 }, 6), call( 'train', { 'learning_rate/model1': 0.018090169943749474, 'learning_rate/model2': 0.009045084971874737, 'momentum/model1': 0.95, 'momentum/model2': 0.9 }, 7), call( 'train', { 'learning_rate/model1': 0.0019098300562505265, 'learning_rate/model2': 0.0009549150281252633, 'momentum/model1': 0.95, 'momentum/model2': 0.9 }, 10) ] else: if by_epoch: calls = [ call('train', { 'learning_rate': 0.018000000000000002, 'momentum': 0.95 }, 1), call('train', { 'learning_rate': 0.02, 'momentum': 0.95 }, 11), call('train', { 'learning_rate': 0.018090169943749474, 'momentum': 0.95 }, 61), call('train', { 'learning_rate': 0.0019098300562505265, 'momentum': 0.95 }, 100) ] else: calls = [ call('train', { 'learning_rate': 0.018000000000000002, 'momentum': 0.95 }, 1), call('train', { 'learning_rate': 0.02, 'momentum': 0.95 }, 6), call('train', { 'learning_rate': 0.018090169943749474, 'momentum': 0.95 }, 7), call('train', { 'learning_rate': 0.0019098300562505265, 'momentum': 0.95 }, 10) ] hook.writer.add_scalars.assert_has_calls(calls, any_order=True)