Spaces:
Sleeping
Sleeping
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved. | |
# Set up custom environment before nearly anything else is imported | |
# NOTE: this should be the first import (no not reorder) | |
from maskrcnn_benchmark.utils.env import setup_environment # noqa F401 isort:skip | |
import argparse | |
import os | |
import torch | |
from maskrcnn_benchmark.config import cfg | |
from maskrcnn_benchmark.data import make_data_loader | |
from maskrcnn_benchmark.engine.inference import inference | |
from maskrcnn_benchmark.modeling.detector import build_detection_model | |
from maskrcnn_benchmark.utils.checkpoint import DetectronCheckpointer | |
from maskrcnn_benchmark.utils.collect_env import collect_env_info | |
from maskrcnn_benchmark.utils.comm import synchronize, get_rank | |
from maskrcnn_benchmark.utils.logger import setup_logger | |
from maskrcnn_benchmark.utils.miscellaneous import mkdir | |
from maskrcnn_benchmark.utils.stats import get_model_complexity_info | |
def run_test(cfg, model, distributed, log_dir): | |
if distributed and hasattr(model, "module"): | |
model = model.module | |
torch.cuda.empty_cache() # TODO check if it helps | |
iou_types = ("bbox",) | |
if cfg.MODEL.MASK_ON: | |
iou_types = iou_types + ("segm",) | |
if cfg.MODEL.KEYPOINT_ON: | |
iou_types = iou_types + ("keypoints",) | |
dataset_names = cfg.DATASETS.TEST | |
if isinstance(dataset_names[0], (list, tuple)): | |
dataset_names = [dataset for group in dataset_names for dataset in group] | |
output_folders = [None] * len(dataset_names) | |
if log_dir: | |
for idx, dataset_name in enumerate(dataset_names): | |
output_folder = os.path.join(log_dir, "inference", dataset_name) | |
mkdir(output_folder) | |
output_folders[idx] = output_folder | |
data_loaders_val = make_data_loader(cfg, is_train=False, is_distributed=distributed) | |
for output_folder, dataset_name, data_loader_val in zip(output_folders, dataset_names, data_loaders_val): | |
inference( | |
model, | |
data_loader_val, | |
dataset_name=dataset_name, | |
iou_types=iou_types, | |
box_only=cfg.MODEL.RPN_ONLY and (cfg.MODEL.RPN_ARCHITECTURE == "RPN" or cfg.DATASETS.CLASS_AGNOSTIC), | |
device=cfg.MODEL.DEVICE, | |
expected_results=cfg.TEST.EXPECTED_RESULTS, | |
expected_results_sigma_tol=cfg.TEST.EXPECTED_RESULTS_SIGMA_TOL, | |
output_folder=output_folder, | |
cfg=cfg, | |
) | |
synchronize() | |
def main(): | |
parser = argparse.ArgumentParser(description="PyTorch Object Detection Inference") | |
parser.add_argument( | |
"--config-file", | |
default="/private/home/fmassa/github/detectron.pytorch_v2/configs/e2e_faster_rcnn_R_50_C4_1x_caffe2.yaml", | |
metavar="FILE", | |
help="path to config file", | |
) | |
parser.add_argument( | |
"--weight", | |
default=None, | |
metavar="FILE", | |
help="path to config file", | |
) | |
parser.add_argument("--local_rank", type=int, default=0) | |
parser.add_argument( | |
"opts", | |
help="Modify config options using the command-line", | |
default=None, | |
nargs=argparse.REMAINDER, | |
) | |
args = parser.parse_args() | |
num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1 | |
distributed = num_gpus > 1 | |
if distributed: | |
torch.cuda.set_device(args.local_rank) | |
torch.distributed.init_process_group(backend="nccl", init_method="env://") | |
cfg.local_rank = args.local_rank | |
cfg.num_gpus = num_gpus | |
cfg.merge_from_file(args.config_file) | |
cfg.merge_from_list(args.opts) | |
cfg.freeze() | |
log_dir = cfg.OUTPUT_DIR | |
if args.weight: | |
log_dir = os.path.join(log_dir, "eval", os.path.splitext(os.path.basename(args.weight))[0]) | |
if log_dir: | |
mkdir(log_dir) | |
logger = setup_logger("maskrcnn_benchmark", log_dir, get_rank()) | |
logger.info(args) | |
logger.info("Using {} GPUs".format(num_gpus)) | |
logger.info(cfg) | |
logger.info("Collecting env info (might take some time)") | |
logger.info("\n" + collect_env_info()) | |
model = build_detection_model(cfg) | |
model.to(cfg.MODEL.DEVICE) | |
params, flops = get_model_complexity_info( | |
model, | |
(3, cfg.INPUT.MAX_SIZE_TEST, cfg.INPUT.MIN_SIZE_TEST), | |
input_constructor=lambda x: {"images": [torch.rand(x).cuda()]}, | |
) | |
print("FLOPs: {}, #Parameter: {}".format(params, flops)) | |
checkpointer = DetectronCheckpointer(cfg, model, save_dir=cfg.OUTPUT_DIR) | |
if args.weight: | |
_ = checkpointer.load(args.weight, force=True) | |
else: | |
_ = checkpointer.load(cfg.MODEL.WEIGHT) | |
run_test(cfg, model, distributed, log_dir) | |
logger.info("FLOPs: {}, #Parameter: {}".format(params, flops)) | |
if __name__ == "__main__": | |
main() | |