Spaces:
Paused
Paused
| # -*- coding: utf-8 -*- | |
| # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved | |
| import logging | |
| import numpy as np | |
| from collections import Counter | |
| import tqdm | |
| from detectron2.checkpoint import DetectionCheckpointer | |
| from detectron2.config import get_cfg | |
| from detectron2.data import build_detection_test_loader | |
| from detectron2.engine import default_argument_parser | |
| from detectron2.modeling import build_model | |
| from detectron2.utils.analysis import ( | |
| activation_count_operators, | |
| flop_count_operators, | |
| parameter_count_table, | |
| ) | |
| from detectron2.utils.logger import setup_logger | |
| logger = logging.getLogger("detectron2") | |
| def setup(args): | |
| cfg = get_cfg() | |
| cfg.merge_from_file(args.config_file) | |
| cfg.DATALOADER.NUM_WORKERS = 0 | |
| cfg.merge_from_list(args.opts) | |
| cfg.freeze() | |
| setup_logger() | |
| return cfg | |
| def do_flop(cfg): | |
| data_loader = build_detection_test_loader(cfg, cfg.DATASETS.TEST[0]) | |
| model = build_model(cfg) | |
| DetectionCheckpointer(model).load(cfg.MODEL.WEIGHTS) | |
| model.eval() | |
| counts = Counter() | |
| total_flops = [] | |
| for idx, data in zip(tqdm.trange(args.num_inputs), data_loader): # noqa | |
| count = flop_count_operators(model, data) | |
| counts += count | |
| total_flops.append(sum(count.values())) | |
| logger.info( | |
| "(G)Flops for Each Type of Operators:\n" + str([(k, v / idx) for k, v in counts.items()]) | |
| ) | |
| logger.info("Total (G)Flops: {}±{}".format(np.mean(total_flops), np.std(total_flops))) | |
| def do_activation(cfg): | |
| data_loader = build_detection_test_loader(cfg, cfg.DATASETS.TEST[0]) | |
| model = build_model(cfg) | |
| DetectionCheckpointer(model).load(cfg.MODEL.WEIGHTS) | |
| model.eval() | |
| counts = Counter() | |
| total_activations = [] | |
| for idx, data in zip(tqdm.trange(args.num_inputs), data_loader): # noqa | |
| count = activation_count_operators(model, data) | |
| counts += count | |
| total_activations.append(sum(count.values())) | |
| logger.info( | |
| "(Million) Activations for Each Type of Operators:\n" | |
| + str([(k, v / idx) for k, v in counts.items()]) | |
| ) | |
| logger.info( | |
| "Total (Million) Activations: {}±{}".format( | |
| np.mean(total_activations), np.std(total_activations) | |
| ) | |
| ) | |
| def do_parameter(cfg): | |
| model = build_model(cfg) | |
| logger.info("Parameter Count:\n" + parameter_count_table(model, max_depth=5)) | |
| def do_structure(cfg): | |
| model = build_model(cfg) | |
| logger.info("Model Structure:\n" + str(model)) | |
| if __name__ == "__main__": | |
| parser = default_argument_parser( | |
| epilog=""" | |
| Examples: | |
| To show parameters of a model: | |
| $ ./analyze_model.py --tasks parameter \\ | |
| --config-file ../configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_1x.yaml | |
| Flops and activations are data-dependent, therefore inputs and model weights | |
| are needed to count them: | |
| $ ./analyze_model.py --num-inputs 100 --tasks flop \\ | |
| --config-file ../configs/COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_1x.yaml \\ | |
| MODEL.WEIGHTS /path/to/model.pkl | |
| """ | |
| ) | |
| parser.add_argument( | |
| "--tasks", | |
| choices=["flop", "activation", "parameter", "structure"], | |
| required=True, | |
| nargs="+", | |
| ) | |
| parser.add_argument( | |
| "--num-inputs", | |
| default=100, | |
| type=int, | |
| help="number of inputs used to compute statistics for flops/activations, " | |
| "both are data dependent.", | |
| ) | |
| args = parser.parse_args() | |
| assert not args.eval_only | |
| assert args.num_gpus == 1 | |
| cfg = setup(args) | |
| for task in args.tasks: | |
| { | |
| "flop": do_flop, | |
| "activation": do_activation, | |
| "parameter": do_parameter, | |
| "structure": do_structure, | |
| }[task](cfg) | |