[WIP] Support LR_ReduceOnPlateau

mmcv/runner/hooks/lr_updater.py

@@ -614,3 +614,184 @@ def format_param(name, optim, param):
         if name not in param:
             raise KeyError(f'{name} is not found in {param.keys()}')
         return param[name]
+
+
+@HOOKS.register_module()
+class ReduceLrUpdateHook(LrUpdaterHook):
+
+    def __init__(self,
+                 periods,
+                 val_metric=None,
+                 mode='min',
+                 factor=0.1,
+                 patience=10,
+                 threshold=1e-4,
+                 threshold_mode='rel',
+                 cooldown=0,
+                 min_lr=0.,
+                 eps=1e-8,
+                 **kwargs):
+        if isinstance(periods, list):
+            assert mmcv.is_list_of(periods, int)
+            assert all([s > 0 for s in periods])
+        else:
+            raise TypeError('"periods" must be a list')
+        self.periods = periods
+        self.val_metric = val_metric
+        if mode not in ['min', 'max']:
+            raise ValueError(
+                'mode must be one of "min" or "max", instead got {mode}')
+        self.mode = mode
+        if factor >= 1.0:
+            raise ValueError('Factor should be < 1.0')
+        self.factor = factor
+        self.patience = patience
+        self.threshold = threshold
+        if threshold_mode not in ['rel', 'abs']:
+            raise ValueError('thresh_mode must be one of "rel" or "abs",\
+                 instead got {threshold_mode}')
+        self.threshold_mode = threshold_mode
+        self.cooldown = cooldown
+        self.cooldown_counter = 0
+        self.best = None
+        self.num_bad_epochs = None
+        self.mode_worse = None  # the worse value for the chosen mode
+        self.min_lr = min_lr
+        self.eps = eps
+        self.last_epoch = 0
+        self._init_is_better(self.mode)
+        self._reset()
+        super(ReduceLrUpdateHook, self).__init__(**kwargs)
+
+    def get_lr(self, runner, regular_lr):
+        if self.num_bad_epochs > self.patience:
+            self.cooldown_counter = self.cooldown
+            self.num_bad_epochs = 0
+            if regular_lr - regular_lr * self.factor > self.eps:
+                new_lr = max(regular_lr * self.factor, self.min_lr)
+            else:
+                new_lr = regular_lr
+            return new_lr
+        else:
+            return regular_lr
+
+    def get_regular_lr(self, runner):
+        if isinstance(runner.optimizer, dict):
+            lr_groups = {}
+            for k in runner.optimizer.keys():
+                _lr_group = [
+                    self.get_lr(runner, _regular_lr)
+                    for _regular_lr in self.regular_lr[k]
+                ]
+                lr_groups.update({k: _lr_group})
+                # self.regular_lr.update({k: _lr_group})
+            return lr_groups
+        else:
+            return [
+                self.get_lr(runner, _regular_lr)
+                for _regular_lr in self.regular_lr
+            ]
+
+    def _init_is_better(self, mode):
+        if mode == 'min':
+            self.mode_worse = float('inf')
+        else:
+            self.mode_worse = float('-inf')
+
+    def _reset(self):
+        self.best = self.mode_worse
+        self.cooldown_counter = 0
+        self.num_bad_epochs = 0
+
+    def is_better(self, a, best):
+        if self.mode == 'min' and self.threshold_mode == 'rel':
+            rel_epsilon = 1. - self.threshold
+            return a < best * rel_epsilon
+
+        elif self.mode == 'min' and self.threshold_mode == 'abs':
+            return a < best - self.threshold
+
+        elif self.mode == 'max' and self.threshold_mode == 'rel':
+            rel_epsilon = 1. + self.threshold
+            return a > best * rel_epsilon
+
+        else:
+            return a > best + self.threshold
+
+    @property
+    def in_cooldown(self):
+        return self.cooldown_counter > 0
+
+    def after_train_epoch(self, runner):
+        if not self.by_epoch:
+            return
+        cur_epoch = runner.epoch
+        if self.warmup is not None and self.warmup_by_epoch:
+            if cur_epoch <= self.warmup_epochs:
+                return
+        if cur_epoch in self.periods and self.val_metric is None:
+            current = runner.outputs.loss
+            if self.is_better(current, self.best):
+                self.best = current
+                self.num_bad_epochs = 0
+            else:
+                self.num_bad_epochs += 1
+
+            if self.in_cooldown:
+                self.cooldown_counter -= 1
+                self.num_bad_epochs = 0
+
+    def after_train_iter(self, runner):
+        if self.by_epoch:
+            return
+        cur_iter = runner.iter
+        if self.warmup_epochs is not None and cur_iter <= self.warmup_iters:
+            return
+        if cur_iter in self.periods and self.val_metric is None:
+            current = runner.outputs.loss
+            if self.is_better(current, self.best):
+                self.best = current
+                self.num_bad_epochs = 0
+            else:
+                self.num_bad_epochs += 1
+
+            if self.in_cooldown:
+                self.cooldown_counter -= 1
+                self.num_bad_epochs = 0
+
+    def after_val_epoch(self, runner):
+        if not self.by_epoch:
+            return
+        cur_epoch = runner.epoch
+        if self.warmup is not None and self.warmup_by_epoch:
+            if cur_epoch <= self.warmup_epochs:
+                return
+        if cur_epoch in self.periods and self.val_metric is not None:
+            current = runner.outputs[self.val_metric]
+            if self.is_better(current, self.best):
+                self.best = current
+                self.num_bad_epochs = 0
+            else:
+                self.num_bad_epochs += 1
+
+            if self.in_cooldown:
+                self.cooldown_counter -= 1
+                self.num_bad_epochs = 0
+
+    def after_val_iter(self, runner):
+        if self.by_epoch:
+            return
+        cur_iter = runner.iter
+        if self.warmup_epochs is not None and cur_iter <= self.warmup_iters:
+            return
+        if cur_iter in self.periods and self.val_metric is not None:
+            current = runner.outputs[self.val_metric]
+            if self.is_better(current, self.best):
+                self.best = current
+                self.num_bad_epochs = 0
+            else:
+                self.num_bad_epochs += 1
+
+            if self.in_cooldown:
+                self.cooldown_counter -= 1
+                self.num_bad_epochs = 0

tests/test_runner/test_hooks.py

@@ -16,6 +16,7 @@
 import torch
 import torch.nn as nn
 from torch.nn.init import constant_
+from torch.optim.lr_scheduler import ReduceLROnPlateau
 from torch.utils.data import DataLoader
 
 from mmcv.runner import (CheckpointHook, EMAHook, IterTimerHook,
@@ -25,6 +26,7 @@
 from mmcv.runner.hooks.lr_updater import (CosineRestartLrUpdaterHook,
                                           CyclicLrUpdaterHook,
                                           OneCycleLrUpdaterHook,
+                                          ReduceLrUpdateHook,
                                           StepLrUpdaterHook)
 
 
@@ -869,6 +871,52 @@ def test_cyclic_lr_update_hook(multi_optimizers, max_iters):
     hook.writer.add_scalars.assert_has_calls(calls, any_order=True)
 
 
+@pytest.mark.parametrize('multi_optimziers', (True, False))
+def test_reduce_lr_update_hook(multi_optimziers):
+    """Test ReduceLrUpdateHook."""
+    with pytest.raises(TypeError):
+        # periods should be specified
+        ReduceLrUpdateHook()
+
+    with pytest.raises(AssertionError):
+        # periods should all be positive
+        ReduceLrUpdateHook(periods=[1, 2, -2])
+
+    with pytest.raises(ValueError):
+        # mode should be either 'min' or 'max'
+        ReduceLrUpdateHook(periods=[0, 1], mode='sum')
+
+    with pytest.raises(ValueError):
+        # factor should be < 1.0
+        ReduceLrUpdateHook(periods=[0, 1], mode='min', factor=1.0)
+
+    with pytest.raises(ValueError):
+        # threshold_mode should be 'rel' or 'abs'
+        ReduceLrUpdateHook(
+            periods=[0, 1], mode='min', factor=1.0, threshold_mode='sum')
+
+    sys.modules['pavi'] = MagicMock()
+    loader = DataLoader(torch.ones((10, 2)))
+    runner = _build_demo_runner(multi_optimziers=multi_optimziers)
+
+    hook = ReduceLROnPlateau(
+        periods=list(range(20)), mode='min', factor=0.1, patience=2)
+    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_optimziers:
+        pass
+    else:
+        pass
+
+
 @pytest.mark.parametrize('log_model', (True, False))
 def test_mlflow_hook(log_model):
     sys.modules['mlflow'] = MagicMock()