[Enhance] Add step momentum updater to support MMDet3D (#1011)

* add StepMomentumUpdaterHook

* add unit test

* fix typos

* refactor step updater

* replace stage with exp

* fix linting error

* use all() operation

mmcv/runner/hooks/lr_updater.py

@@ -2,6 +2,7 @@
 import numbers
 from math import cos, pi
 
+import mmcv
 from .hook import HOOKS, Hook
 
 
@@ -177,10 +178,9 @@ class StepLrUpdaterHook(LrUpdaterHook):
     """
 
     def __init__(self, step, gamma=0.1, min_lr=None, **kwargs):
-        assert isinstance(step, (list, int))
         if isinstance(step, list):
-            for s in step:
-                assert isinstance(s, int) and s > 0
+            assert mmcv.is_list_of(step, int)
+            assert all([s > 0 for s in step])
         elif isinstance(step, int):
             assert step > 0
         else:
@@ -193,19 +193,17 @@ class StepLrUpdaterHook(LrUpdaterHook):
     def get_lr(self, runner, base_lr):
         progress = runner.epoch if self.by_epoch else runner.iter
 
+        # calculate exponential term
         if isinstance(self.step, int):
-            lr = base_lr * (self.gamma**(progress // self.step))
-            if self.min_lr is not None:
-                # clip to a minimum value
-                lr = max(lr, self.min_lr)
-            return lr
-
-        exp = len(self.step)
-        for i, s in enumerate(self.step):
-            if progress < s:
-                exp = i
-                break
-        lr = base_lr * self.gamma**exp
+            exp = progress // self.step
+        else:
+            exp = len(self.step)
+            for i, s in enumerate(self.step):
+                if progress < s:
+                    exp = i
+                    break
+
+        lr = base_lr * (self.gamma**exp)
         if self.min_lr is not None:
             # clip to a minimum value
             lr = max(lr, self.min_lr)

mmcv/runner/hooks/momentum_updater.py

+import mmcv
 from .hook import HOOKS, Hook
 from .lr_updater import annealing_cos, annealing_linear, format_param
 
@@ -148,6 +149,54 @@ class MomentumUpdaterHook(Hook):
                 self._set_momentum(runner, warmup_momentum)
 
 
+@HOOKS.register_module()
+class StepMomentumUpdaterHook(MomentumUpdaterHook):
+    """Step momentum scheduler with min value clipping.
+
+    Args:
+        step (int | list[int]): Step to decay the momentum. If an int value is
+            given, regard it as the decay interval. If a list is given, decay
+            momentum at these steps.
+        gamma (float, optional): Decay momentum ratio. Default: 0.5.
+        min_momentum (float, optional): Minimum momentum value to keep. If
+            momentum after decay is lower than this value, it will be clipped
+            accordingly. If None is given, we don't perform lr clipping.
+            Default: None.
+    """
+
+    def __init__(self, step, gamma=0.5, min_momentum=None, **kwargs):
+        if isinstance(step, list):
+            assert mmcv.is_list_of(step, int)
+            assert all([s > 0 for s in step])
+        elif isinstance(step, int):
+            assert step > 0
+        else:
+            raise TypeError('"step" must be a list or integer')
+        self.step = step
+        self.gamma = gamma
+        self.min_momentum = min_momentum
+        super(StepMomentumUpdaterHook, self).__init__(**kwargs)
+
+    def get_momentum(self, runner, base_momentum):
+        progress = runner.epoch if self.by_epoch else runner.iter
+
+        # calculate exponential term
+        if isinstance(self.step, int):
+            exp = progress // self.step
+        else:
+            exp = len(self.step)
+            for i, s in enumerate(self.step):
+                if progress < s:
+                    exp = i
+                    break
+
+        momentum = base_momentum * (self.gamma**exp)
+        if self.min_momentum is not None:
+            # clip to a minimum value
+            momentum = max(momentum, self.min_momentum)
+        return momentum
+
+
 @HOOKS.register_module()
 class CosineAnnealingMomentumUpdaterHook(MomentumUpdaterHook):
 
@@ -419,11 +468,12 @@ class OneCycleMomentumUpdaterHook(MomentumUpdaterHook):
             end_iter = phase['end_iter']
             if curr_iter <= end_iter or i == len(self.momentum_phases) - 1:
                 pct = (curr_iter - start_iter) / (end_iter - start_iter)
-                lr = self.anneal_func(param_group[phase['start_momentum']],
-                                      param_group[phase['end_momentum']], pct)
+                momentum = self.anneal_func(
+                    param_group[phase['start_momentum']],
+                    param_group[phase['end_momentum']], pct)
                 break
             start_iter = end_iter
-        return lr
+        return momentum
 
     def get_regular_momentum(self, runner):
         if isinstance(runner.optimizer, dict):

tests/test_runner/test_hooks.py

@@ -540,7 +540,7 @@ def test_cosine_restart_lr_update_hook(multi_optimziers):
 
 
 @pytest.mark.parametrize('multi_optimziers', (True, False))
-def test_step_lr_update_hook(multi_optimziers):
+def test_step_runner_hook(multi_optimziers):
     """Test StepLrUpdaterHook."""
     with pytest.raises(TypeError):
         # `step` should be specified
@@ -557,6 +557,15 @@ def test_step_lr_update_hook(multi_optimziers):
     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)
@@ -583,36 +592,36 @@ def test_step_lr_update_hook(multi_optimziers):
                 'train', {
                     'learning_rate/model1': 0.01,
                     'learning_rate/model2': 0.005,
-                    'momentum/model1': 0.95,
-                    'momentum/model2': 0.9
+                    'momentum/model1': 0.475,
+                    'momentum/model2': 0.45
                 }, 6),
             call(
                 'train', {
                     'learning_rate/model1': 0.0025,
                     'learning_rate/model2': 0.00125,
-                    'momentum/model1': 0.95,
-                    'momentum/model2': 0.9
+                    'momentum/model1': 0.11875,
+                    'momentum/model2': 0.1125
                 }, 16),
             call(
                 'train', {
                     'learning_rate/model1': 0.00125,
                     'learning_rate/model2': 0.001,
-                    'momentum/model1': 0.95,
-                    'momentum/model2': 0.9
+                    'momentum/model1': 0.059375,
+                    'momentum/model2': 0.05625
                 }, 21),
             call(
                 'train', {
                     'learning_rate/model1': 0.001,
                     'learning_rate/model2': 0.001,
-                    'momentum/model1': 0.95,
-                    'momentum/model2': 0.9
+                    'momentum/model1': 0.05,
+                    'momentum/model2': 0.05
                 }, 26),
             call(
                 'train', {
                     'learning_rate/model1': 0.001,
                     'learning_rate/model2': 0.001,
-                    'momentum/model1': 0.95,
-                    'momentum/model2': 0.9
+                    'momentum/model1': 0.05,
+                    'momentum/model2': 0.05
                 }, 30)
         ]
     else:
@@ -623,23 +632,23 @@ def test_step_lr_update_hook(multi_optimziers):
             }, 1),
             call('train', {
                 'learning_rate': 0.01,
-                'momentum': 0.95
+                'momentum': 0.475
             }, 6),
             call('train', {
                 'learning_rate': 0.0025,
-                'momentum': 0.95
+                'momentum': 0.11875
             }, 16),
             call('train', {
                 'learning_rate': 0.00125,
-                'momentum': 0.95
+                'momentum': 0.059375
             }, 21),
             call('train', {
                 'learning_rate': 0.001,
-                'momentum': 0.95
+                'momentum': 0.05
             }, 26),
             call('train', {
                 'learning_rate': 0.001,
-                'momentum': 0.95
+                'momentum': 0.05
             }, 30)
         ]
     hook.writer.add_scalars.assert_has_calls(calls, any_order=True)
@@ -649,6 +658,14 @@ def test_step_lr_update_hook(multi_optimziers):
     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)
@@ -675,22 +692,22 @@ def test_step_lr_update_hook(multi_optimziers):
                 'train', {
                     'learning_rate/model1': 0.002,
                     'learning_rate/model2': 0.001,
-                    'momentum/model1': 0.95,
-                    'momentum/model2': 0.9
+                    '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': 0.95,
-                    'momentum/model2': 0.9
+                    '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': 0.95,
-                    'momentum/model2': 0.9
+                    'momentum/model1': 9.500000000000002e-4,
+                    'momentum/model2': 9.000000000000002e-4
                 }, 9)
         ]
     else:
@@ -701,16 +718,18 @@ def test_step_lr_update_hook(multi_optimziers):
             }, 1),
             call('train', {
                 'learning_rate': 0.002,
-                'momentum': 0.95
+                'momentum': 0.095
             }, 5),
-            call('train', {
-                'learning_rate': 2.0000000000000004e-4,
-                'momentum': 0.95
-            }, 7),
-            call('train', {
-                'learning_rate': 2.0000000000000005e-05,
-                'momentum': 0.95
-            }, 9)
+            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)