diff --git a/docs/zh_cn/tutorials/registry.md b/docs/zh_cn/tutorials/registry.md index 3c209c6207f2aaab2b08221e7c2dd2f4446ef16c..164cb0ea2b9dba8ba16b150289b0c88406179b7c 100644 --- a/docs/zh_cn/tutorials/registry.md +++ b/docs/zh_cn/tutorials/registry.md @@ -222,7 +222,9 @@ MMEngine 的注册器支æŒè·¨é¡¹ç›®è°ƒç”¨ï¼Œå³å¯ä»¥åœ¨ä¸€ä¸ªé¡¹ç›®ä¸ä½¿ç”¨ - WEIGHT_INITIALIZERS: æƒé‡åˆå§‹åŒ–的工具 - OPTIMIZERS: 注册了 PyTorch ä¸æ‰€æœ‰çš„ `optimizer` 以åŠè‡ªå®šä¹‰çš„ `optimizer` - OPTIMIZER_CONSTRUCTORS: optimizer çš„æž„é€ å™¨ +- PARAM_SCHEDULERS: å„ç§å‚数调度器, 如 `MultiStepLR` - TASK_UTILS: 任务强相关的一些组件,如 `AnchorGenerator`, `BboxCoder` +- EVALUATORS: 用于验è¯æ¨¡åž‹ç²¾åº¦çš„评估器 下é¢æˆ‘们以 OpenMMLab å¼€æºé¡¹ç›®ä¸ºä¾‹ä»‹ç»å¦‚何跨项目调用模å—。 diff --git a/mmengine/evaluator/__init__.py b/mmengine/evaluator/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c2a8d5dda8238e5c600f6c60ec30805743f9fde7 --- /dev/null +++ b/mmengine/evaluator/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) OpenMMLab. All rights reserved. +from .base import BaseEvaluator +from .builder import build_evaluator +from .composed_evaluator import ComposedEvaluator + +__all__ = ['BaseEvaluator', 'ComposedEvaluator', 'build_evaluator'] diff --git a/mmengine/evaluator/base.py b/mmengine/evaluator/base.py new file mode 100644 index 0000000000000000000000000000000000000000..287c2fe272822e668730a929df48a153c0bbd270 --- /dev/null +++ b/mmengine/evaluator/base.py @@ -0,0 +1,210 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import os.path as osp +import pickle +import shutil +import tempfile +import warnings +from abc import ABCMeta, abstractmethod +from typing import Any, List, Optional, Union + +import torch +import torch.distributed as dist + +from mmengine.utils import mkdir_or_exist + + +class BaseEvaluator(metaclass=ABCMeta): + """Base class for an evaluator. + + The evaluator first processes each batch of data_samples and + predictions, and appends the processed results in to the results list. + Then it collects all results together from all ranks if distributed + training is used. Finally, it computes the metrics of the entire dataset. + + Args: + collect_device (str): Device name used for collecting results from + different ranks during distributed training. Must be 'cpu' or + 'gpu'. Defaults to 'cpu'. + """ + + def __init__(self, collect_device: str = 'cpu') -> None: + self._dataset_meta: Union[None, dict] = None + self.collect_device = collect_device + self.results: List[Any] = [] + + rank, world_size = get_dist_info() + self.rank = rank + self.world_size = world_size + + @property + def dataset_meta(self) -> Optional[dict]: + return self._dataset_meta + + @dataset_meta.setter + def dataset_meta(self, dataset_meta: dict) -> None: + self._dataset_meta = dataset_meta + + @abstractmethod + def process(self, data_samples: dict, predictions: dict) -> None: + """Process one batch of data samples and predictions. The processed + results should be stored in ``self.results``, which will be used to + compute the metrics when all batches have been processed. + + Args: + data_samples (dict): The data samples from the dataset. + predictions (dict): The output of the model. + """ + + @abstractmethod + def compute_metrics(self, results: list) -> dict: + """Compute the metrics from processed results. + + Args: + results (list): The processed results of each batch. + Returns: + dict: The computed metrics. The keys are the names of the metrics, + and the values are corresponding results. + """ + + def evaluate(self, size: int) -> dict: + """Evaluate the model performance of the whole dataset after processing + all batches. + + Args: + size (int): Length of the entire validation dataset. When batch + size > 1, the dataloader may pad some data samples to make + sure all ranks have the same length of dataset slice. The + ``collect_results`` function will drop the padded data base on + this size. + + Returns: + metrics (dict): Evaluation metrics dict on the val dataset. The + keys are the names of the metrics, and the values are + corresponding results. + """ + if len(self.results) == 0: + warnings.warn( + f'{self.__class__.__name__} got empty `self._results`. Please ' + 'ensure that the processed results are properly added into ' + '`self._results` in `process` method.') + + if self.world_size == 1: + # non-distributed + results = self.results + else: + results = collect_results(self.results, size, self.collect_device) + + if self.rank == 0: + # TODO: replace with mmengine.dist.master_only + metrics = [self.compute_metrics(results)] + else: + metrics = [None] # type: ignore + # TODO: replace with mmengine.dist.broadcast + if self.world_size > 1: + metrics = dist.broadcast_object_list(metrics) + + # reset the results list + self.results.clear() + return metrics[0] + + +# TODO: replace with mmengine.dist.get_dist_info +def get_dist_info(): + if dist.is_available() and dist.is_initialized(): + rank = dist.get_rank() + world_size = dist.get_world_size() + else: + rank = 0 + world_size = 1 + return rank, world_size + + +# TODO: replace with mmengine.dist.collect_results +def collect_results(results, size, device='cpu'): + """Collected results in distributed environments.""" + # TODO: replace with mmengine.dist.collect_results + if device == 'gpu': + return collect_results_gpu(results, size) + elif device == 'cpu': + return collect_results_cpu(results, size) + else: + NotImplementedError(f"device must be 'cpu' or 'gpu', but got {device}") + + +# TODO: replace with mmengine.dist.collect_results +def collect_results_cpu(result_part, size, tmpdir=None): + rank, world_size = get_dist_info() + # create a tmp dir if it is not specified + if tmpdir is None: + MAX_LEN = 512 + # 32 is whitespace + dir_tensor = torch.full((MAX_LEN, ), + 32, + dtype=torch.uint8, + device='cuda') + if rank == 0: + mkdir_or_exist('.dist_test') + tmpdir = tempfile.mkdtemp(dir='.dist_test') + tmpdir = torch.tensor( + bytearray(tmpdir.encode()), dtype=torch.uint8, device='cuda') + dir_tensor[:len(tmpdir)] = tmpdir + dist.broadcast(dir_tensor, 0) + tmpdir = dir_tensor.cpu().numpy().tobytes().decode().rstrip() + else: + mkdir_or_exist(tmpdir) + # dump the part result to the dir + with open(osp.join(tmpdir, f'part_{rank}.pkl'), 'wb') as f: + pickle.dump(result_part, f, protocol=2) + dist.barrier() + # collect all parts + if rank != 0: + return None + else: + # load results of all parts from tmp dir + part_list = [] + for i in range(world_size): + with open(osp.join(tmpdir, f'part_{i}.pkl'), 'wb') as f: + part_list.append(pickle.load(f)) + # sort the results + ordered_results = [] + for res in zip(*part_list): + ordered_results.extend(list(res)) + # the dataloader may pad some samples + ordered_results = ordered_results[:size] + # remove tmp dir + shutil.rmtree(tmpdir) + return ordered_results + + +# TODO: replace with mmengine.dist.collect_results +def collect_results_gpu(result_part, size): + rank, world_size = get_dist_info() + # dump result part to tensor with pickle + part_tensor = torch.tensor( + bytearray(pickle.dumps(result_part)), dtype=torch.uint8, device='cuda') + # gather all result part tensor shape + shape_tensor = torch.tensor(part_tensor.shape, device='cuda') + shape_list = [shape_tensor.clone() for _ in range(world_size)] + dist.all_gather(shape_list, shape_tensor) + # padding result part tensor to max length + shape_max = torch.tensor(shape_list).max() + part_send = torch.zeros(shape_max, dtype=torch.uint8, device='cuda') + part_send[:shape_tensor[0]] = part_tensor + part_recv_list = [ + part_tensor.new_zeros(shape_max) for _ in range(world_size) + ] + # gather all result part + dist.all_gather(part_recv_list, part_send) + + if rank == 0: + part_list = [] + for recv, shape in zip(part_recv_list, shape_list): + part_list.append( + pickle.loads(recv[:shape[0]].cpu().numpy().tobytes())) + # sort the results + ordered_results = [] + for res in zip(*part_list): + ordered_results.extend(list(res)) + # the dataloader may pad some samples + ordered_results = ordered_results[:size] + return ordered_results diff --git a/mmengine/evaluator/builder.py b/mmengine/evaluator/builder.py new file mode 100644 index 0000000000000000000000000000000000000000..710c6554b96e63b94c5237a1c31168b40e3bfb23 --- /dev/null +++ b/mmengine/evaluator/builder.py @@ -0,0 +1,16 @@ +# Copyright (c) OpenMMLab. All rights reserved. +from ..registry import EVALUATORS +from .composed_evaluator import ComposedEvaluator + + +def build_evaluator(cfg: dict) -> object: + """Build function of evaluator. + + When the evaluator config is a list, it will automatically build composed + evaluators. + """ + if isinstance(cfg, list): + evaluators = [EVALUATORS.build(_cfg) for _cfg in cfg] + return ComposedEvaluator(evaluators=evaluators) + else: + return EVALUATORS.build(cfg) diff --git a/mmengine/evaluator/composed_evaluator.py b/mmengine/evaluator/composed_evaluator.py new file mode 100644 index 0000000000000000000000000000000000000000..225284e70a43208949df0391487b2f53d45e74d9 --- /dev/null +++ b/mmengine/evaluator/composed_evaluator.py @@ -0,0 +1,73 @@ +# Copyright (c) OpenMMLab. All rights reserved. +from typing import Optional, Sequence, Union + +from .base import BaseEvaluator + + +class ComposedEvaluator: + """Wrapper class to compose multiple :class:`DatasetEvaluator` instances. + + Args: + evaluators (Sequence[BaseEvaluator]): The evaluators to compose. + collect_device (str): Device name used for collecting results from + different ranks during distributed training. Must be 'cpu' or + 'gpu'. Defaults to 'cpu'. + """ + + def __init__(self, + evaluators: Sequence[BaseEvaluator], + collect_device='cpu'): + self._dataset_meta: Union[None, dict] = None + self.collect_device = collect_device + self.evaluators = evaluators + + @property + def dataset_meta(self) -> Optional[dict]: + return self._dataset_meta + + @dataset_meta.setter + def dataset_meta(self, dataset_meta: dict) -> None: + self._dataset_meta = dataset_meta + for evaluator in self.evaluators: + evaluator.dataset_meta = dataset_meta + + def process(self, data_samples: dict, predictions: dict): + """Invoke process method of each wrapped evaluator. + + Args: + data_samples (dict): The data samples from the dataset. + predictions (dict): The output of the model. + """ + + for evalutor in self.evaluators: + evalutor.process(data_samples, predictions) + + def evaluate(self, size: int) -> dict: + """Invoke evaluate method of each wrapped evaluator and collect the + metrics dict. + + Args: + size (int): Length of the entire validation dataset. When batch + size > 1, the dataloader may pad some data samples to make + sure all ranks have the same length of dataset slice. The + ``collect_results`` function will drop the padded data base on + this size. + + Returns: + metrics (dict): Evaluation metrics of all wrapped evaluators. The + keys are the names of the metrics, and the values are + corresponding results. + """ + metrics = {} + for evaluator in self.evaluators: + _metrics = evaluator.evaluate(size) + + # Check metric name conflicts + for name in _metrics.keys(): + if name in metrics: + raise ValueError( + 'There are multiple evaluators with the same metric ' + f'name {name}') + + metrics.update(_metrics) + return metrics diff --git a/mmengine/registry/__init__.py b/mmengine/registry/__init__.py index d8602d4f5f38b4863c44bab4819e814c98363f69..24ebcb6e7f3546bb0e4b61e8c54b0701d52d052b 100644 --- a/mmengine/registry/__init__.py +++ b/mmengine/registry/__init__.py @@ -1,6 +1,6 @@ # Copyright (c) OpenMMLab. All rights reserved. from .registry import Registry, build_from_cfg -from .root import (DATA_SAMPLERS, DATASETS, HOOKS, MODELS, +from .root import (DATA_SAMPLERS, DATASETS, EVALUATORS, HOOKS, MODELS, OPTIMIZER_CONSTRUCTORS, OPTIMIZERS, PARAM_SCHEDULERS, RUNNER_CONSTRUCTORS, RUNNERS, TASK_UTILS, TRANSFORMS, WEIGHT_INITIALIZERS) @@ -8,5 +8,6 @@ from .root import (DATA_SAMPLERS, DATASETS, HOOKS, MODELS, __all__ = [ 'Registry', 'build_from_cfg', 'RUNNERS', 'RUNNER_CONSTRUCTORS', 'HOOKS', 'DATASETS', 'DATA_SAMPLERS', 'TRANSFORMS', 'MODELS', 'WEIGHT_INITIALIZERS', - 'OPTIMIZERS', 'OPTIMIZER_CONSTRUCTORS', 'TASK_UTILS', 'PARAM_SCHEDULERS' + 'OPTIMIZERS', 'OPTIMIZER_CONSTRUCTORS', 'TASK_UTILS', 'PARAM_SCHEDULERS', + 'EVALUATORS' ] diff --git a/mmengine/registry/root.py b/mmengine/registry/root.py index c67a8d5b67f294aad22e915f7d8d171e750fa866..9636fae21f135b0d94124c53c3aecf5a814a3acf 100644 --- a/mmengine/registry/root.py +++ b/mmengine/registry/root.py @@ -34,3 +34,6 @@ PARAM_SCHEDULERS = Registry('parameter scheduler') # manage task-specific modules like anchor generators and box coders TASK_UTILS = Registry('task util') + +# manage all kinds of evaluators for computing metrics +EVALUATORS = Registry('evaluator')