Use MMCV's EvalHook in MMDetection (#4806)

* Use EvalHook in MMCV

* Add DeprecationWarning

* update

* fix unit test

* add comment for unit test

Author: hhaAndroid

mmdet/core/evaluation/eval_hooks.py

@@ -1,250 +1,30 @@
 import os.path as osp
-import warnings
-from math import inf
 
-import mmcv
 import torch.distributed as dist
-from mmcv.runner import Hook
+from mmcv.runner import DistEvalHook as BaseDistEvalHook
+from mmcv.runner import EvalHook as BaseEvalHook
 from torch.nn.modules.batchnorm import _BatchNorm
-from torch.utils.data import DataLoader
 
-from mmdet.utils import get_root_logger
 
+class EvalHook(BaseEvalHook):
 
-class EvalHook(Hook):
-    """Evaluation hook.
-
-    Notes:
-        If new arguments are added for EvalHook, tools/test.py,
-        tools/analysis_tools/eval_metric.py may be effected.
-
-    Attributes:
-        dataloader (DataLoader): A PyTorch dataloader.
-        start (int, optional): Evaluation starting epoch. It enables evaluation
-            before the training starts if ``start`` <= the resuming epoch.
-            If None, whether to evaluate is merely decided by ``interval``.
-            Default: None.
-        interval (int): Evaluation interval (by epochs). Default: 1.
-        save_best (str, optional): If a metric is specified, it would measure
-            the best checkpoint during evaluation. The information about best
-            checkpoint would be save in best.json.
-            Options are the evaluation metrics to the test dataset. e.g.,
-            ``bbox_mAP``, ``segm_mAP`` for bbox detection and instance
-            segmentation. ``AR@100`` for proposal recall. If ``save_best`` is
-            ``auto``, the first key will be used. The interval of
-            ``CheckpointHook`` should device EvalHook. Default: None.
-        rule (str, optional): Comparison rule for best score. If set to None,
-            it will infer a reasonable rule. Keys such as 'mAP' or 'AR' will
-            be inferred by 'greater' rule. Keys contain 'loss' will be inferred
-             by 'less' rule. Options are 'greater', 'less'. Default: None.
-        **eval_kwargs: Evaluation arguments fed into the evaluate function of
-            the dataset.
-    """
-
-    rule_map = {'greater': lambda x, y: x > y, 'less': lambda x, y: x < y}
-    init_value_map = {'greater': -inf, 'less': inf}
-    greater_keys = ['mAP', 'AR']
-    less_keys = ['loss']
-
-    def __init__(self,
-                 dataloader,
-                 start=None,
-                 interval=1,
-                 by_epoch=True,
-                 save_best=None,
-                 rule=None,
-                 **eval_kwargs):
-        if not isinstance(dataloader, DataLoader):
-            raise TypeError('dataloader must be a pytorch DataLoader, but got'
-                            f' {type(dataloader)}')
-        if not interval > 0:
-            raise ValueError(f'interval must be positive, but got {interval}')
-        if start is not None and start < 0:
-            warnings.warn(
-                f'The evaluation start epoch {start} is smaller than 0, '
-                f'use 0 instead', UserWarning)
-            start = 0
-        self.dataloader = dataloader
-        self.interval = interval
-        self.by_epoch = by_epoch
-        self.start = start
-        assert isinstance(save_best, str) or save_best is None
-        self.save_best = save_best
-        self.eval_kwargs = eval_kwargs
-        self.initial_epoch_flag = True
-
-        self.logger = get_root_logger()
-
-        if self.save_best is not None:
-            self._init_rule(rule, self.save_best)
-
-    def _init_rule(self, rule, key_indicator):
-        """Initialize rule, key_indicator, comparison_func, and best score.
-
-        Args:
-            rule (str | None): Comparison rule for best score.
-            key_indicator (str | None): Key indicator to determine the
-                comparison rule.
-        """
-        if rule not in self.rule_map and rule is not None:
-            raise KeyError(f'rule must be greater, less or None, '
-                           f'but got {rule}.')
-
-        if rule is None:
-            if key_indicator != 'auto':
-                if any(key in key_indicator for key in self.greater_keys):
-                    rule = 'greater'
-                elif any(key in key_indicator for key in self.less_keys):
-                    rule = 'less'
-                else:
-                    raise ValueError(f'Cannot infer the rule for key '
-                                     f'{key_indicator}, thus a specific rule '
-                                     f'must be specified.')
-        self.rule = rule
-        self.key_indicator = key_indicator
-        if self.rule is not None:
-            self.compare_func = self.rule_map[self.rule]
-
-    def before_run(self, runner):
-        if self.save_best is not None:
-            if runner.meta is None:
-                warnings.warn('runner.meta is None. Creating a empty one.')
-                runner.meta = dict()
-            runner.meta.setdefault('hook_msgs', dict())
-
-    def before_train_epoch(self, runner):
-        """Evaluate the model only at the start of training."""
-        if not self.initial_epoch_flag:
+    def _do_evaluate(self, runner):
+        """perform evaluation and save ckpt."""
+        if not self._should_evaluate(runner):
             return
-        if self.start is not None and runner.epoch >= self.start:
-            self.after_train_epoch(runner)
-        self.initial_epoch_flag = False
-
-    def evaluation_flag(self, runner):
-        """Judge whether to perform_evaluation after this epoch.
 
-        Returns:
-            bool: The flag indicating whether to perform evaluation.
-        """
-        if self.start is None:
-            if not self.every_n_epochs(runner, self.interval):
-                # No evaluation during the interval epochs.
-                return False
-        elif (runner.epoch + 1) < self.start:
-            # No evaluation if start is larger than the current epoch.
-            return False
-        else:
-            # Evaluation only at epochs 3, 5, 7... if start==3 and interval==2
-            if (runner.epoch + 1 - self.start) % self.interval:
-                return False
-        return True
-
-    def after_train_epoch(self, runner):
-        if not self.by_epoch or not self.evaluation_flag(runner):
-            return
         from mmdet.apis import single_gpu_test
         results = single_gpu_test(runner.model, self.dataloader, show=False)
+        runner.log_buffer.output['eval_iter_num'] = len(self.dataloader)
         key_score = self.evaluate(runner, results)
         if self.save_best:
-            self.save_best_checkpoint(runner, key_score)
-
-    def after_train_iter(self, runner):
-        if self.by_epoch or not self.every_n_iters(runner, self.interval):
-            return
-        from mmdet.apis import single_gpu_test
-        results = single_gpu_test(runner.model, self.dataloader, show=False)
-        key_score = self.evaluate(runner, results)
-        if self.save_best:
-            self.save_best_checkpoint(runner, key_score)
-
-    def save_best_checkpoint(self, runner, key_score):
-        best_score = runner.meta['hook_msgs'].get(
-            'best_score', self.init_value_map[self.rule])
-        if self.compare_func(key_score, best_score):
-            best_score = key_score
-            runner.meta['hook_msgs']['best_score'] = best_score
-            last_ckpt = runner.meta['hook_msgs']['last_ckpt']
-            runner.meta['hook_msgs']['best_ckpt'] = last_ckpt
-            mmcv.symlink(
-                last_ckpt,
-                osp.join(runner.work_dir, f'best_{self.key_indicator}.pth'))
-            time_stamp = runner.epoch + 1 if self.by_epoch else runner.iter + 1
-            self.logger.info(f'Now best checkpoint is epoch_{time_stamp}.pth.'
-                             f'Best {self.key_indicator} is {best_score:0.4f}')
-
-    def evaluate(self, runner, results):
-        eval_res = self.dataloader.dataset.evaluate(
-            results, logger=runner.logger, **self.eval_kwargs)
-        for name, val in eval_res.items():
-            runner.log_buffer.output[name] = val
-        runner.log_buffer.ready = True
-        if self.save_best is not None:
-            if self.key_indicator == 'auto':
-                # infer from eval_results
-                self._init_rule(self.rule, list(eval_res.keys())[0])
-            return eval_res[self.key_indicator]
-        else:
-            return None
-
+            self._save_ckpt(runner, key_score)
 
-class DistEvalHook(EvalHook):
-    """Distributed evaluation hook.
 
-    Notes:
-        If new arguments are added, tools/test.py may be effected.
+class DistEvalHook(BaseDistEvalHook):
 
-    Attributes:
-        dataloader (DataLoader): A PyTorch dataloader.
-        start (int, optional): Evaluation starting epoch. It enables evaluation
-            before the training starts if ``start`` <= the resuming epoch.
-            If None, whether to evaluate is merely decided by ``interval``.
-            Default: None.
-        interval (int): Evaluation interval (by epochs). Default: 1.
-        tmpdir (str | None): Temporary directory to save the results of all
-            processes. Default: None.
-        gpu_collect (bool): Whether to use gpu or cpu to collect results.
-            Default: False.
-        save_best (str, optional): If a metric is specified, it would measure
-            the best checkpoint during evaluation. The information about best
-            checkpoint would be save in best.json.
-            Options are the evaluation metrics to the test dataset. e.g.,
-            ``bbox_mAP``, ``segm_mAP`` for bbox detection and instance
-            segmentation. ``AR@100`` for proposal recall. If ``save_best`` is
-            ``auto``, the first key will be used. The interval of
-            ``CheckpointHook`` should device EvalHook. Default: None.
-        rule (str | None): Comparison rule for best score. If set to None,
-            it will infer a reasonable rule. Default: 'None'.
-        broadcast_bn_buffer (bool): Whether to broadcast the
-            buffer(running_mean and running_var) of rank 0 to other rank
-            before evaluation. Default: True.
-        **eval_kwargs: Evaluation arguments fed into the evaluate function of
-            the dataset.
-    """
-
-    def __init__(self,
-                 dataloader,
-                 start=None,
-                 interval=1,
-                 by_epoch=True,
-                 tmpdir=None,
-                 gpu_collect=False,
-                 save_best=None,
-                 rule=None,
-                 broadcast_bn_buffer=True,
-                 **eval_kwargs):
-        super().__init__(
-            dataloader,
-            start=start,
-            interval=interval,
-            by_epoch=by_epoch,
-            save_best=save_best,
-            rule=rule,
-            **eval_kwargs)
-        self.broadcast_bn_buffer = broadcast_bn_buffer
-        self.tmpdir = tmpdir
-        self.gpu_collect = gpu_collect
-
-    def _broadcast_bn_buffer(self, runner):
+    def _do_evaluate(self, runner):
+        """perform evaluation and save ckpt."""
         # Synchronization of BatchNorm's buffer (running_mean
         # and running_var) is not supported in the DDP of pytorch,
         # which may cause the inconsistent performance of models in
@@ -258,46 +38,23 @@ def _broadcast_bn_buffer(self, runner):
                     dist.broadcast(module.running_var, 0)
                     dist.broadcast(module.running_mean, 0)
 
-    def after_train_epoch(self, runner):
-        if not self.by_epoch or not self.evaluation_flag(runner):
+        if not self._should_evaluate(runner):
             return
 
-        if self.broadcast_bn_buffer:
-            self._broadcast_bn_buffer(runner)
-
-        from mmdet.apis import multi_gpu_test
         tmpdir = self.tmpdir
         if tmpdir is None:
             tmpdir = osp.join(runner.work_dir, '.eval_hook')
-        results = multi_gpu_test(
-            runner.model,
-            self.dataloader,
-            tmpdir=tmpdir,
-            gpu_collect=self.gpu_collect)
-        if runner.rank == 0:
-            print('\n')
-            key_score = self.evaluate(runner, results)
-            if self.save_best:
-                self.save_best_checkpoint(runner, key_score)
-
-    def after_train_iter(self, runner):
-        if self.by_epoch or not self.every_n_iters(runner, self.interval):
-            return
-
-        if self.broadcast_bn_buffer:
-            self._broadcast_bn_buffer(runner)
 
         from mmdet.apis import multi_gpu_test
-        tmpdir = self.tmpdir
-        if tmpdir is None:
-            tmpdir = osp.join(runner.work_dir, '.eval_hook')
         results = multi_gpu_test(
             runner.model,
             self.dataloader,
             tmpdir=tmpdir,
             gpu_collect=self.gpu_collect)
         if runner.rank == 0:
             print('\n')
+            runner.log_buffer.output['eval_iter_num'] = len(self.dataloader)
             key_score = self.evaluate(runner, results)
+
             if self.save_best:
-                self.save_best_checkpoint(runner, key_score)
+                self._save_ckpt(runner, key_score)

tests/test_data/test_datasets/test_common.py

@@ -274,9 +274,12 @@ def test_evaluation_hook(EvalHookParam):
     runner.run([dataloader], [('train', 1)], 2)
     assert evalhook.evaluate.call_count == 3  # before epoch1 and after e1 & e2
 
+    # the evaluation start epoch cannot be less than 0
     runner = _build_demo_runner()
-    with pytest.warns(UserWarning):
-        evalhook = EvalHookParam(dataloader, start=-2)
+    with pytest.raises(ValueError):
+        EvalHookParam(dataloader, start=-2)
+
+    evalhook = EvalHookParam(dataloader, start=0)
     evalhook.evaluate = MagicMock()
     runner.register_hook(evalhook)
     runner.run([dataloader], [('train', 1)], 2)