128 lines
3.6 KiB
Python
128 lines
3.6 KiB
Python
# -*- 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)
|