diff --git a/mmcv/runner/hooks/lr_updater.py b/mmcv/runner/hooks/lr_updater.py index 4aa22ed1c0..2edf5610b0 100644 --- a/mmcv/runner/hooks/lr_updater.py +++ b/mmcv/runner/hooks/lr_updater.py @@ -367,13 +367,16 @@ class CyclicLrUpdaterHook(LrUpdaterHook): triangular policy inside a cycle. This improves the performance in the 3D detection area. - Attributes: + Args: + by_epoch (bool): Whether to update LR by epoch. target_ratio (tuple[float]): Relative ratio of the highest LR and the lowest LR to the initial LR. cyclic_times (int): Number of cycles during training step_ratio_up (float): The ratio of the increasing process of LR in the total cycle. - by_epoch (bool): Whether to update LR by epoch. + anneal_strategy (str): {'cos', 'linear'} + Specifies the annealing strategy: 'cos' for cosine annealing, + 'linear' for linear annealing. Default: 'cos'. """ def __init__(self, @@ -381,6 +384,7 @@ def __init__(self, target_ratio=(10, 1e-4), cyclic_times=1, step_ratio_up=0.4, + anneal_strategy='cos', **kwargs): if isinstance(target_ratio, float): target_ratio = (target_ratio, target_ratio / 1e5) @@ -400,6 +404,14 @@ def __init__(self, self.cyclic_times = cyclic_times self.step_ratio_up = step_ratio_up self.lr_phases = [] # init lr_phases + # validate anneal_strategy + if anneal_strategy not in ['cos', 'linear']: + raise ValueError('anneal_strategy must be one of "cos" or ' + f'"linear", instead got {anneal_strategy}') + elif anneal_strategy == 'cos': + self.anneal_func = annealing_cos + elif anneal_strategy == 'linear': + self.anneal_func = annealing_linear assert not by_epoch, \ 'currently only support "by_epoch" = False' @@ -425,9 +437,9 @@ def get_lr(self, runner, base_lr): curr_iter %= max_iter_per_phase if start_iter <= curr_iter < end_iter: progress = curr_iter - start_iter - return annealing_cos(base_lr * start_ratio, - base_lr * end_ratio, - progress / (end_iter - start_iter)) + return self.anneal_func(base_lr * start_ratio, + base_lr * end_ratio, + progress / (end_iter - start_iter)) @HOOKS.register_module() diff --git a/tests/test_runner/test_hooks.py b/tests/test_runner/test_hooks.py index 6a6e72960d..41ce493b14 100644 --- a/tests/test_runner/test_hooks.py +++ b/tests/test_runner/test_hooks.py @@ -22,6 +22,7 @@ MlflowLoggerHook, PaviLoggerHook, WandbLoggerHook, build_runner) from mmcv.runner.hooks.lr_updater import (CosineRestartLrUpdaterHook, + CyclicLrUpdaterHook, OneCycleLrUpdaterHook, StepLrUpdaterHook) @@ -714,6 +715,93 @@ def test_step_lr_update_hook(multi_optimziers): hook.writer.add_scalars.assert_has_calls(calls, any_order=True) +@pytest.mark.parametrize('multi_optimizers, max_iters', [(True, 8), + (False, 8)]) +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) + + @pytest.mark.parametrize('log_model', (True, False)) def test_mlflow_hook(log_model): sys.modules['mlflow'] = MagicMock()