Spaces:
Runtime error
Runtime error
| # Copyright (c) OpenMMLab. All rights reserved. | |
| import argparse | |
| import tempfile | |
| from pathlib import Path | |
| import torch | |
| from mmengine import Config, DictAction | |
| from mmengine.logging import MMLogger | |
| from mmengine.model import revert_sync_batchnorm | |
| from mmengine.registry import init_default_scope | |
| from mmseg.models import BaseSegmentor | |
| from mmseg.registry import MODELS | |
| from mmseg.structures import SegDataSample | |
| import os | |
| import json | |
| try: | |
| from mmengine.analysis import get_model_complexity_info | |
| from mmengine.analysis.print_helper import _format_size | |
| except ImportError: | |
| raise ImportError('Please upgrade mmengine >= 0.6.0 to use this script.') | |
| from fvcore.nn import FlopCountAnalysis | |
| def parse_args(): | |
| parser = argparse.ArgumentParser( | |
| description='Get the FLOPs of a segmentor') | |
| parser.add_argument('config', help='train config file path') | |
| parser.add_argument( | |
| '--shape', | |
| type=int, | |
| nargs='+', | |
| default=[512, 512], | |
| help='input image size') | |
| 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.') | |
| args = parser.parse_args() | |
| return args | |
| def inference(args: argparse.Namespace, logger: MMLogger) -> dict: | |
| config_name = Path(args.config) | |
| if not config_name.exists(): | |
| logger.error(f'Config file {config_name} does not exist') | |
| cfg: Config = Config.fromfile(config_name) | |
| cfg.work_dir = tempfile.TemporaryDirectory().name | |
| cfg.log_level = 'WARN' | |
| if args.cfg_options is not None: | |
| cfg.merge_from_dict(args.cfg_options) | |
| init_default_scope(cfg.get('scope', 'mmseg')) | |
| if len(args.shape) == 1: | |
| input_shape = (3, args.shape[0], args.shape[0]) | |
| elif len(args.shape) == 2: | |
| input_shape = (3, ) + tuple(args.shape) | |
| else: | |
| raise ValueError('invalid input shape') | |
| result = {} | |
| model: BaseSegmentor = MODELS.build(cfg.model) | |
| if hasattr(model, 'auxiliary_head'): | |
| model.auxiliary_head = None | |
| if torch.cuda.is_available(): | |
| model.cuda() | |
| model = revert_sync_batchnorm(model) | |
| result['ori_shape'] = input_shape[-2:] | |
| result['pad_shape'] = input_shape[-2:] | |
| data_batch = { | |
| 'inputs': [torch.rand(input_shape)], | |
| 'data_samples': [SegDataSample(metainfo=result)] | |
| } | |
| data = model.data_preprocessor(data_batch) | |
| model.eval() | |
| if cfg.model.decode_head.type in ['MaskFormerHead', 'Mask2FormerHead']: | |
| # TODO: Support MaskFormer and Mask2Former | |
| raise NotImplementedError('MaskFormer and Mask2Former are not ' | |
| 'supported yet.') | |
| if hasattr(model, 'module'): | |
| all_cfgs = model.module.backbone.all_cfgs | |
| else: | |
| all_cfgs = model.backbone.all_cfgs | |
| stitch_results = {} | |
| for cfg_id in all_cfgs: | |
| if hasattr(model, 'module'): | |
| model.module.backbone.reset_stitch_id(cfg_id) | |
| else: | |
| model.backbone.reset_stitch_id(cfg_id) | |
| flops = FlopCountAnalysis(model, torch.randn([1]+list(input_shape)).cuda()).total() | |
| stitch_results[cfg_id] = flops | |
| save_dir = './model_flops' | |
| if not os.path.exists(save_dir): | |
| os.mkdir(save_dir) | |
| config_name = args.config.split('/')[-1].split('.')[0] | |
| with open(os.path.join(save_dir, f'snnet_flops_{config_name}.json'), 'w+') as f: | |
| json.dump(stitch_results, f, indent=4) | |
| def main(): | |
| args = parse_args() | |
| logger = MMLogger.get_instance(name='MMLogger') | |
| inference(args, logger) | |
| if __name__ == '__main__': | |
| main() | |