[Feature] Support for linearly learning rate decay (#1627)

* update support for linearly learning rate decay

* Fix LinearAnnealingLrUpdaterHook, update LinearAnnealingMomentumUpdaterHook, add unit test

add docstring

add docstring

update linear lr momentum schedule test

fix ci

Fix CI

* add lr and momentum hook to runner and hooks package

add lr and momentum hook to runner and hooks package

* replace multi_optimziers with multi_optimizers
Co-authored-by: HAOCHENYE <21724054@zju.edu.cn>
Co-authored-by: Mashiro <57566630+HAOCHENYE@users.noreply.github.com>

docs/en/community/contributing.md

-../../../CONTRIBUTING.md
\ No newline at end of file
+../../../CONTRIBUTING.md

mmcv/runner/__init__.py

@@ -23,10 +23,12 @@ from .hooks.lr_updater import (CosineAnnealingLrUpdaterHook,
                                CosineRestartLrUpdaterHook, CyclicLrUpdaterHook,
                                ExpLrUpdaterHook, FixedLrUpdaterHook,
                                FlatCosineAnnealingLrUpdaterHook,
-                               InvLrUpdaterHook, LrUpdaterHook,
-                               OneCycleLrUpdaterHook, PolyLrUpdaterHook)
+                               InvLrUpdaterHook, LinearAnnealingLrUpdaterHook,
+                               LrUpdaterHook, OneCycleLrUpdaterHook,
+                               PolyLrUpdaterHook)
 from .hooks.momentum_updater import (CosineAnnealingMomentumUpdaterHook,
                                      CyclicMomentumUpdaterHook,
+                                     LinearAnnealingMomentumUpdaterHook,
                                      MomentumUpdaterHook,
                                      OneCycleMomentumUpdaterHook,
                                      StepMomentumUpdaterHook)
@@ -62,5 +64,6 @@ __all__ = [
     '_load_checkpoint_with_prefix', 'EvalHook', 'DistEvalHook', 'Sequential',
     'ModuleDict', 'ModuleList', 'GradientCumulativeOptimizerHook',
     'GradientCumulativeFp16OptimizerHook', 'DefaultRunnerConstructor',
-    'SegmindLoggerHook'
+    'SegmindLoggerHook', 'LinearAnnealingMomentumUpdaterHook',
+    'LinearAnnealingLrUpdaterHook'
 ]

mmcv/runner/hooks/__init__.py

@@ -12,11 +12,14 @@ from .lr_updater import (CosineAnnealingLrUpdaterHook,
                          CosineRestartLrUpdaterHook, CyclicLrUpdaterHook,
                          ExpLrUpdaterHook, FixedLrUpdaterHook,
                          FlatCosineAnnealingLrUpdaterHook, InvLrUpdaterHook,
-                         LrUpdaterHook, OneCycleLrUpdaterHook,
-                         PolyLrUpdaterHook, StepLrUpdaterHook)
+                         LinearAnnealingLrUpdaterHook, LrUpdaterHook,
+                         OneCycleLrUpdaterHook, PolyLrUpdaterHook,
+                         StepLrUpdaterHook)
 from .memory import EmptyCacheHook
 from .momentum_updater import (CosineAnnealingMomentumUpdaterHook,
-                               CyclicMomentumUpdaterHook, MomentumUpdaterHook,
+                               CyclicMomentumUpdaterHook,
+                               LinearAnnealingMomentumUpdaterHook,
+                               MomentumUpdaterHook,
                                OneCycleMomentumUpdaterHook,
                                StepMomentumUpdaterHook)
 from .optimizer import (Fp16OptimizerHook, GradientCumulativeFp16OptimizerHook,
@@ -39,5 +42,6 @@ __all__ = [
     'CyclicMomentumUpdaterHook', 'OneCycleMomentumUpdaterHook',
     'SyncBuffersHook', 'EMAHook', 'EvalHook', 'DistEvalHook', 'ProfilerHook',
     'GradientCumulativeOptimizerHook', 'GradientCumulativeFp16OptimizerHook',
-    'SegmindLoggerHook'
+    'SegmindLoggerHook', 'LinearAnnealingLrUpdaterHook',
+    'LinearAnnealingMomentumUpdaterHook'
 ]

mmcv/runner/hooks/lr_updater.py

@@ -256,6 +256,14 @@ class InvLrUpdaterHook(LrUpdaterHook):
 
 @HOOKS.register_module()
 class CosineAnnealingLrUpdaterHook(LrUpdaterHook):
+    """CosineAnnealing LR scheduler.
+
+    Args:
+        min_lr (float, optional): The minimum lr. Default: None.
+        min_lr_ratio (float, optional): The ratio of minimum lr to the base lr.
+            Either `min_lr` or `min_lr_ratio` should be specified.
+            Default: None.
+    """
 
     def __init__(self, min_lr=None, min_lr_ratio=None, **kwargs):
         assert (min_lr is None) ^ (min_lr_ratio is None)
@@ -644,6 +652,38 @@ class OneCycleLrUpdaterHook(LrUpdaterHook):
         return lr
 
 
+@HOOKS.register_module()
+class LinearAnnealingLrUpdaterHook(LrUpdaterHook):
+    """Linear annealing LR Scheduler decays the learning rate of each parameter
+    group linearly.
+
+    Args:
+        min_lr (float, optional): The minimum lr. Default: None.
+        min_lr_ratio (float, optional): The ratio of minimum lr to the base lr.
+            Either `min_lr` or `min_lr_ratio` should be specified.
+            Default: None.
+    """
+
+    def __init__(self, min_lr=None, min_lr_ratio=None, **kwargs):
+        assert (min_lr is None) ^ (min_lr_ratio is None)
+        self.min_lr = min_lr
+        self.min_lr_ratio = min_lr_ratio
+        super(LinearAnnealingLrUpdaterHook, self).__init__(**kwargs)
+
+    def get_lr(self, runner, base_lr):
+        if self.by_epoch:
+            progress = runner.epoch
+            max_progress = runner.max_epochs
+        else:
+            progress = runner.iter
+            max_progress = runner.max_iters
+        if self.min_lr_ratio is not None:
+            target_lr = base_lr * self.min_lr_ratio
+        else:
+            target_lr = self.min_lr
+        return annealing_linear(base_lr, target_lr, progress / max_progress)
+
+
 def annealing_cos(start, end, factor, weight=1):
     """Calculate annealing cos learning rate.
 

mmcv/runner/hooks/momentum_updater.py

@@ -200,6 +200,15 @@ class StepMomentumUpdaterHook(MomentumUpdaterHook):
 
 @HOOKS.register_module()
 class CosineAnnealingMomentumUpdaterHook(MomentumUpdaterHook):
+    """Cosine annealing LR Momentum decays the Momentum of each parameter group
+    linearly.
+
+    Args:
+        min_momentum (float, optional): The minimum momentum. Default: None.
+        min_momentum_ratio (float, optional): The ratio of minimum momentum to
+            the base momentum. Either `min_momentum` or `min_momentum_ratio`
+            should be specified. Default: None.
+    """
 
     def __init__(self, min_momentum=None, min_momentum_ratio=None, **kwargs):
         assert (min_momentum is None) ^ (min_momentum_ratio is None)
@@ -222,6 +231,39 @@ class CosineAnnealingMomentumUpdaterHook(MomentumUpdaterHook):
                              progress / max_progress)
 
 
+@HOOKS.register_module()
+class LinearAnnealingMomentumUpdaterHook(MomentumUpdaterHook):
+    """Linear annealing LR Momentum decays the Momentum of each parameter group
+    linearly.
+
+    Args:
+        min_momentum (float, optional): The minimum momentum. Default: None.
+        min_momentum_ratio (float, optional): The ratio of minimum momentum to
+            the base momentum. Either `min_momentum` or `min_momentum_ratio`
+            should be specified. Default: None.
+    """
+
+    def __init__(self, min_momentum=None, min_momentum_ratio=None, **kwargs):
+        assert (min_momentum is None) ^ (min_momentum_ratio is None)
+        self.min_momentum = min_momentum
+        self.min_momentum_ratio = min_momentum_ratio
+        super(LinearAnnealingMomentumUpdaterHook, self).__init__(**kwargs)
+
+    def get_momentum(self, runner, base_momentum):
+        if self.by_epoch:
+            progress = runner.epoch
+            max_progress = runner.max_epochs
+        else:
+            progress = runner.iter
+            max_progress = runner.max_iters
+        if self.min_momentum_ratio is not None:
+            target_momentum = base_momentum * self.min_momentum_ratio
+        else:
+            target_momentum = self.min_momentum
+        return annealing_linear(base_momentum, target_momentum,
+                                progress / max_progress)
+
+
 @HOOKS.register_module()
 class CyclicMomentumUpdaterHook(MomentumUpdaterHook):
     """Cyclic momentum Scheduler.

tests/test_runner/test_hooks.py

@@ -711,6 +711,83 @@ def test_cosine_runner_hook(multi_optimizers):
     hook.writer.add_scalars.assert_has_calls(calls, any_order=True)
 
 
+@pytest.mark.parametrize('multi_optimizers', (True, False))
+def test_linear_runner_hook(multi_optimizers):
+    sys.modules['pavi'] = MagicMock()
+    loader = DataLoader(torch.ones((10, 2)))
+    runner = _build_demo_runner(multi_optimizers=multi_optimizers)
+
+    # add momentum scheduler
+
+    hook_cfg = dict(
+        type='LinearAnnealingMomentumUpdaterHook',
+        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='LinearAnnealingLrUpdaterHook',
+        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')
+    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.01,
+                    'learning_rate/model2': 0.005,
+                    'momentum/model1': 0.97,
+                    'momentum/model2': 0.9189473684210527,
+                }, 6),
+            call(
+                'train', {
+                    'learning_rate/model1': 0.0019999999999999983,
+                    'learning_rate/model2': 0.0009999999999999992,
+                    'momentum/model1': 0.9860000000000001,
+                    'momentum/model2': 0.9341052631578949,
+                }, 10)
+        ]
+    else:
+        calls = [
+            call('train', {
+                'learning_rate': 0.02,
+                'momentum': 0.95
+            }, 1),
+            call('train', {
+                'learning_rate': 0.01,
+                'momentum': 0.97
+            }, 6),
+            call(
+                'train', {
+                    'learning_rate': 0.0019999999999999983,
+                    'momentum': 0.9860000000000001
+                }, 10)
+        ]
+    hook.writer.add_scalars.assert_has_calls(calls, any_order=True)
+
+
 @pytest.mark.parametrize('multi_optimizers, by_epoch', [(False, False),
                                                         (True, False),
                                                         (False, True),