[Feature]: Support calculating FLOPs of detectors (#9777)

Author: zwhus

Diff for: .dev_scripts/benchmark_valid_flops.py

@@ -0,0 +1,295 @@
+import logging
+import re
+import tempfile
+from argparse import ArgumentParser
+from collections import OrderedDict
+from functools import partial
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import torch
+from mmengine import Config, DictAction
+from mmengine.analysis import get_model_complexity_info
+from mmengine.analysis.print_helper import _format_size
+from mmengine.fileio import FileClient
+from mmengine.logging import MMLogger
+from mmengine.model import revert_sync_batchnorm
+from mmengine.runner import Runner
+from modelindex.load_model_index import load
+from rich.console import Console
+from rich.table import Table
+from rich.text import Text
+from tqdm import tqdm
+
+from mmdet.registry import MODELS
+from mmdet.utils import register_all_modules
+
+console = Console()
+MMDET_ROOT = Path(__file__).absolute().parents[1]
+
+
+def parse_args():
+    parser = ArgumentParser(description='Valid all models in model-index.yml')
+    parser.add_argument(
+        '--shape',
+        type=int,
+        nargs='+',
+        default=[1280, 800],
+        help='input image size')
+    parser.add_argument(
+        '--checkpoint_root',
+        help='Checkpoint file root path. If set, load checkpoint before test.')
+    parser.add_argument('--img', default='demo/demo.jpg', help='Image file')
+    parser.add_argument('--models', nargs='+', help='models name to inference')
+    parser.add_argument(
+        '--batch-size',
+        type=int,
+        default=1,
+        help='The batch size during the inference.')
+    parser.add_argument(
+        '--flops', action='store_true', help='Get Flops and Params of models')
+    parser.add_argument(
+        '--flops-str',
+        action='store_true',
+        help='Output FLOPs and params counts in a string form.')
+    parser.add_argument(
+        '--cfg-options',
+        nargs='+',
+        action=DictAction,
+        help='override some settings in the used config, the key-value pair '
+        'in xxx=yyy format will be merged into config file. If the value to '
+        'be overwritten is a list, it should be like key="[a,b]" or key=a,b '
+        'It also allows nested list/tuple values, e.g. key="[(a,b),(c,d)]" '
+        'Note that the quotation marks are necessary and that no white space '
+        'is allowed.')
+    parser.add_argument(
+        '--size_divisor',
+        type=int,
+        default=32,
+        help='Pad the input image, the minimum size that is divisible '
+        'by size_divisor, -1 means do not pad the image.')
+    args = parser.parse_args()
+    return args
+
+
+def inference(config_file, checkpoint, work_dir, args, exp_name):
+    logger = MMLogger.get_instance(name='MMLogger')
+    logger.warning('if you want test flops, please make sure torch>=1.12')
+    cfg = Config.fromfile(config_file)
+    cfg.work_dir = work_dir
+    cfg.load_from = checkpoint
+    cfg.log_level = 'WARN'
+    cfg.experiment_name = exp_name
+    if args.cfg_options is not None:
+        cfg.merge_from_dict(args.cfg_options)
+
+    # forward the model
+    result = {'model': config_file.stem}
+
+    if args.flops:
+
+        if len(args.shape) == 1:
+            h = w = args.shape[0]
+        elif len(args.shape) == 2:
+            h, w = args.shape
+        else:
+            raise ValueError('invalid input shape')
+        divisor = args.size_divisor
+        if divisor > 0:
+            h = int(np.ceil(h / divisor)) * divisor
+            w = int(np.ceil(w / divisor)) * divisor
+
+        input_shape = (3, h, w)
+        result['resolution'] = input_shape
+
+        try:
+            cfg = Config.fromfile(config_file)
+            if hasattr(cfg, 'head_norm_cfg'):
+                cfg['head_norm_cfg'] = dict(type='SyncBN', requires_grad=True)
+                cfg['model']['roi_head']['bbox_head']['norm_cfg'] = dict(
+                    type='SyncBN', requires_grad=True)
+                cfg['model']['roi_head']['mask_head']['norm_cfg'] = dict(
+                    type='SyncBN', requires_grad=True)
+
+            if args.cfg_options is not None:
+                cfg.merge_from_dict(args.cfg_options)
+
+            model = MODELS.build(cfg.model)
+            input = torch.rand(1, *input_shape)
+            if torch.cuda.is_available():
+                model.cuda()
+                input = input.cuda()
+            model = revert_sync_batchnorm(model)
+            inputs = (input, )
+            model.eval()
+            outputs = get_model_complexity_info(
+                model, input_shape, inputs, show_table=False, show_arch=False)
+            flops = outputs['flops']
+            params = outputs['params']
+            activations = outputs['activations']
+            result['Get Types'] = 'direct'
+        except:  # noqa 772
+            logger = MMLogger.get_instance(name='MMLogger')
+            logger.warning(
+                'Direct get flops failed, try to get flops with data')
+            cfg = Config.fromfile(config_file)
+            if hasattr(cfg, 'head_norm_cfg'):
+                cfg['head_norm_cfg'] = dict(type='SyncBN', requires_grad=True)
+                cfg['model']['roi_head']['bbox_head']['norm_cfg'] = dict(
+                    type='SyncBN', requires_grad=True)
+                cfg['model']['roi_head']['mask_head']['norm_cfg'] = dict(
+                    type='SyncBN', requires_grad=True)
+            data_loader = Runner.build_dataloader(cfg.val_dataloader)
+            data_batch = next(iter(data_loader))
+            model = MODELS.build(cfg.model)
+            if torch.cuda.is_available():
+                model = model.cuda()
+            model = revert_sync_batchnorm(model)
+            model.eval()
+            _forward = model.forward
+            data = model.data_preprocessor(data_batch)
+            del data_loader
+            model.forward = partial(
+                _forward, data_samples=data['data_samples'])
+            outputs = get_model_complexity_info(
+                model,
+                input_shape,
+                data['inputs'],
+                show_table=False,
+                show_arch=False)
+            flops = outputs['flops']
+            params = outputs['params']
+            activations = outputs['activations']
+            result['Get Types'] = 'dataloader'
+
+        if args.flops_str:
+            flops = _format_size(flops)
+            params = _format_size(params)
+            activations = _format_size(activations)
+
+        result['flops'] = flops
+        result['params'] = params
+
+    return result
+
+
+def show_summary(summary_data, args):
+    table = Table(title='Validation Benchmark Regression Summary')
+    table.add_column('Model')
+    table.add_column('Validation')
+    table.add_column('Resolution (c, h, w)')
+    if args.flops:
+        table.add_column('Flops', justify='right', width=11)
+        table.add_column('Params', justify='right')
+
+    for model_name, summary in summary_data.items():
+        row = [model_name]
+        valid = summary['valid']
+        color = 'green' if valid == 'PASS' else 'red'
+        row.append(f'[{color}]{valid}[/{color}]')
+        if valid == 'PASS':
+            row.append(str(summary['resolution']))
+            if args.flops:
+                row.append(str(summary['flops']))
+                row.append(str(summary['params']))
+        table.add_row(*row)
+
+    console.print(table)
+    table_data = {
+        x.header: [Text.from_markup(y).plain for y in x.cells]
+        for x in table.columns
+    }
+    table_pd = pd.DataFrame(table_data)
+    table_pd.to_csv('./mmdetection_flops.csv')
+
+
+# Sample test whether the inference code is correct
+def main(args):
+    register_all_modules()
+    model_index_file = MMDET_ROOT / 'model-index.yml'
+    model_index = load(str(model_index_file))
+    model_index.build_models_with_collections()
+    models = OrderedDict({model.name: model for model in model_index.models})
+
+    logger = MMLogger(
+        'validation',
+        logger_name='validation',
+        log_file='benchmark_test_image.log',
+        log_level=logging.INFO)
+
+    if args.models:
+        patterns = [
+            re.compile(pattern.replace('+', '_')) for pattern in args.models
+        ]
+        filter_models = {}
+        for k, v in models.items():
+            k = k.replace('+', '_')
+            if any([re.match(pattern, k) for pattern in patterns]):
+                filter_models[k] = v
+        if len(filter_models) == 0:
+            print('No model found, please specify models in:')
+            print('\n'.join(models.keys()))
+            return
+        models = filter_models
+
+    summary_data = {}
+    tmpdir = tempfile.TemporaryDirectory()
+    for model_name, model_info in tqdm(models.items()):
+
+        if model_info.config is None:
+            continue
+
+        model_info.config = model_info.config.replace('%2B', '+')
+        config = Path(model_info.config)
+
+        try:
+            config.exists()
+        except:  # noqa 722
+            logger.error(f'{model_name}: {config} not found.')
+            continue
+
+        logger.info(f'Processing: {model_name}')
+
+        http_prefix = 'https://download.openmmlab.com/mmdetection/'
+        if args.checkpoint_root is not None:
+            root = args.checkpoint_root
+            if 's3://' in args.checkpoint_root:
+                from petrel_client.common.exception import AccessDeniedError
+                file_client = FileClient.infer_client(uri=root)
+                checkpoint = file_client.join_path(
+                    root, model_info.weights[len(http_prefix):])
+                try:
+                    exists = file_client.exists(checkpoint)
+                except AccessDeniedError:
+                    exists = False
+            else:
+                checkpoint = Path(root) / model_info.weights[len(http_prefix):]
+                exists = checkpoint.exists()
+            if exists:
+                checkpoint = str(checkpoint)
+            else:
+                print(f'WARNING: {model_name}: {checkpoint} not found.')
+                checkpoint = None
+        else:
+            checkpoint = None
+
+        try:
+            # build the model from a config file and a checkpoint file
+            result = inference(MMDET_ROOT / config, checkpoint, tmpdir.name,
+                               args, model_name)
+            result['valid'] = 'PASS'
+        except Exception:  # noqa 722
+            import traceback
+            logger.error(f'"{config}" :\n{traceback.format_exc()}')
+            result = {'valid': 'FAIL'}
+
+        summary_data[model_name] = result
+
+    tmpdir.cleanup()
+    show_summary(summary_data, args)
+
+
+if __name__ == '__main__':
+    args = parse_args()
+    main(args)

Diff for: configs/simple_copy_paste/metafile.yml

@@ -25,7 +25,7 @@ Collections:
 Models:
   - Name: mask-rcnn_r50_fpn_syncbn-all_rpn-2conv_ssj_32x2_270k_coco
     In Collection: SimpleCopyPaste
-    Config: configs/simplecopypaste/mask-rcnn_r50_fpn_rpn-2conv_4conv1fc_syncbn-all_32xb2-ssj-270k_coco.py
+    Config: configs/simple_copy_paste/mask-rcnn_r50_fpn_rpn-2conv_4conv1fc_syncbn-all_32xb2-ssj-270k_coco.py
     Metadata:
       Training Memory (GB): 7.2
       Iterations: 270000
@@ -42,7 +42,7 @@ Models:
 
   - Name: mask-rcnn_r50_fpn_syncbn-all_rpn-2conv_ssj_32x2_90k_coco
     In Collection: SimpleCopyPaste
-    Config: configs/simplecopypaste/mask-rcnn_r50_fpn_rpn-2conv_4conv1fc_syncbn-all_32xb2-ssj-90k_coco.py
+    Config: configs/simple_copy_paste/mask-rcnn_r50_fpn_rpn-2conv_4conv1fc_syncbn-all_32xb2-ssj-90k_coco.py
     Metadata:
       Training Memory (GB): 7.2
       Iterations: 90000
@@ -59,7 +59,7 @@ Models:
 
   - Name: mask-rcnn_r50_fpn_syncbn-all_rpn-2conv_ssj_scp_32x2_270k_coco
     In Collection: SimpleCopyPaste
-    Config: configs/simplecopypaste/mask-rcnn_r50_fpn_rpn-2conv_4conv1fc_syncbn-all_32xb2-ssj-scp-270k_coco.py
+    Config: configs/simple_copy_paste/mask-rcnn_r50_fpn_rpn-2conv_4conv1fc_syncbn-all_32xb2-ssj-scp-270k_coco.py
     Metadata:
       Training Memory (GB): 7.2
       Iterations: 270000
@@ -76,7 +76,7 @@ Models:
 
   - Name: mask-rcnn_r50_fpn_syncbn-all_rpn-2conv_ssj_scp_32x2_90k_coco
     In Collection: SimpleCopyPaste
-    Config: configs/simplecopypaste/mask-rcnn_r50_fpn_rpn-2conv_4conv1fc_syncbn-all_32xb2-ssj-scp-90k_coco.py
+    Config: configs/simple_copy_paste/mask-rcnn_r50_fpn_rpn-2conv_4conv1fc_syncbn-all_32xb2-ssj-scp-90k_coco.py
     Metadata:
       Training Memory (GB): 7.2
       Iterations: 90000

Diff for: mmdet/models/roi_heads/sparse_roi_head.py

@@ -578,7 +578,11 @@ def forward(self, x: Tuple[Tensor], rpn_results_list: InstanceList,
                     batch_img_metas=batch_img_metas,
                     batch_gt_instances=batch_gt_instances)
                 bbox_results.pop('loss_bbox')
-                all_stage_bbox_results.append((bbox_results, ))
+                # torch.jit does not support obj:SamplingResult
+                bbox_results.pop('results_list')
+                bbox_res = bbox_results.copy()
+                bbox_res.pop('sampling_results')
+                all_stage_bbox_results.append((bbox_res, ))
 
                 if self.with_mask:
                     attn_feats = bbox_results['attn_feats']