| |
| from mmcv.cnn import VGG |
| from mmengine.hooks import Hook |
| from mmengine.runner import Runner |
|
|
| from mmdet.registry import HOOKS |
|
|
|
|
| @HOOKS.register_module() |
| class NumClassCheckHook(Hook): |
| """Check whether the `num_classes` in head matches the length of `classes` |
| in `dataset.metainfo`.""" |
|
|
| def _check_head(self, runner: Runner, mode: str) -> None: |
| """Check whether the `num_classes` in head matches the length of |
| `classes` in `dataset.metainfo`. |
| |
| Args: |
| runner (:obj:`Runner`): The runner of the training or evaluation |
| process. |
| """ |
| assert mode in ['train', 'val'] |
| model = runner.model |
| dataset = runner.train_dataloader.dataset if mode == 'train' else \ |
| runner.val_dataloader.dataset |
| if dataset.metainfo.get('classes', None) is None: |
| runner.logger.warning( |
| f'Please set `classes` ' |
| f'in the {dataset.__class__.__name__} `metainfo` and' |
| f'check if it is consistent with the `num_classes` ' |
| f'of head') |
| else: |
| classes = dataset.metainfo['classes'] |
| assert type(classes) is not str, \ |
| (f'`classes` in {dataset.__class__.__name__}' |
| f'should be a tuple of str.' |
| f'Add comma if number of classes is 1 as ' |
| f'classes = ({classes},)') |
| from mmdet.models.roi_heads.mask_heads import FusedSemanticHead |
| for name, module in model.named_modules(): |
| if hasattr(module, 'num_classes') and not name.endswith( |
| 'rpn_head') and not isinstance( |
| module, (VGG, FusedSemanticHead)): |
| assert module.num_classes == len(classes), \ |
| (f'The `num_classes` ({module.num_classes}) in ' |
| f'{module.__class__.__name__} of ' |
| f'{model.__class__.__name__} does not matches ' |
| f'the length of `classes` ' |
| f'{len(classes)}) in ' |
| f'{dataset.__class__.__name__}') |
|
|
| def before_train_epoch(self, runner: Runner) -> None: |
| """Check whether the training dataset is compatible with head. |
| |
| Args: |
| runner (:obj:`Runner`): The runner of the training or evaluation |
| process. |
| """ |
| self._check_head(runner, 'train') |
|
|
| def before_val_epoch(self, runner: Runner) -> None: |
| """Check whether the dataset in val epoch is compatible with head. |
| |
| Args: |
| runner (:obj:`Runner`): The runner of the training or evaluation |
| process. |
| """ |
| self._check_head(runner, 'val') |
|
|
