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.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, is_main_process | |
| 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 | |
| import os | |
| import functools | |
| import io | |
| import os | |
| import datetime | |
| import wandb | |
| import torch | |
| import torch.distributed as dist | |
| import pdb | |
| from pprint import pprint | |
| def init_distributed_mode(args): | |
| """Initialize distributed training, if appropriate""" | |
| if "RANK" in os.environ and "WORLD_SIZE" in os.environ: | |
| args.rank = int(os.environ["RANK"]) | |
| args.world_size = int(os.environ["WORLD_SIZE"]) | |
| args.gpu = int(os.environ["LOCAL_RANK"]) | |
| elif "SLURM_PROCID" in os.environ: | |
| args.rank = int(os.environ["SLURM_PROCID"]) | |
| args.gpu = args.rank % torch.cuda.device_count() | |
| else: | |
| print("Not using distributed mode") | |
| args.distributed = False | |
| return | |
| # args.distributed = True | |
| torch.cuda.set_device(args.gpu) | |
| args.dist_backend = "nccl" | |
| print("| distributed init (rank {}): {}".format(args.rank, args.dist_url), flush=True) | |
| dist.init_process_group( | |
| backend=args.dist_backend, | |
| init_method=args.dist_url, | |
| world_size=args.world_size, | |
| rank=args.rank, | |
| timeout=datetime.timedelta(0, 72000), | |
| ) | |
| dist.barrier() | |
| setup_for_distributed(args.rank == 0) | |
| def setup_for_distributed(is_master): | |
| """ | |
| This function disables printing when not in master process | |
| """ | |
| import builtins as __builtin__ | |
| builtin_print = __builtin__.print | |
| def print(*args, **kwargs): | |
| force = kwargs.pop("force", False) | |
| if is_master or force: | |
| builtin_print(*args, **kwargs) | |
| __builtin__.print = print | |
| def main(): | |
| parser = argparse.ArgumentParser(description="PyTorch Detection to Grounding Inference") | |
| parser.add_argument( | |
| "--config-file", | |
| default="configs/grounding/e2e_dyhead_SwinT_S_FPN_1x_od_grounding_eval.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 | |
| ) | |
| parser.add_argument("--world-size", default=1, type=int, help="number of distributed processes") | |
| parser.add_argument("--dist-url", default="env://", help="url used to set up distributed training") | |
| parser.add_argument("--task_config", default=None) | |
| parser.add_argument("--eval_negative", action="store_true") | |
| parser.add_argument("--wandb_project_name", default="haroldli/language_det_eval") | |
| 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://" | |
| # ) | |
| init_distributed_mode(args) | |
| print("Passed distributed init") | |
| 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) | |
| # we currently disable this | |
| # 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) | |
| if args.weight: | |
| weight_iter = os.path.splitext(os.path.basename(args.weight))[0].split("_")[-1] | |
| try: | |
| weight_iter = int(weight_iter) | |
| except: | |
| weight_iter = 1 | |
| else: | |
| weight_iter = 1 | |
| # get the wandb name | |
| train_wandb_name = os.path.basename(cfg.OUTPUT_DIR) | |
| eval_wandb_name = train_wandb_name + "_eval" + "_Fixed{}_Chunk{}".format(not cfg.DATASETS.LVIS_USE_NORMAL_AP, cfg.TEST.CHUNKED_EVALUATION) | |
| if args.eval_negative: | |
| from maskrcnn_benchmark.engine.inference_contrastive import inference | |
| inference_function = inference | |
| else: | |
| from maskrcnn_benchmark.engine.inference import inference | |
| inference_function = inference | |
| if is_main_process() and train_wandb_name != "__test__": | |
| api = wandb.Api() | |
| runs = api.runs(args.wandb_project_name) | |
| matched_run = None | |
| history = [] | |
| exclude_keys = ['_runtime', '_timestamp'] | |
| for run in runs: | |
| if run.name == eval_wandb_name and str(run._state) == "finished": | |
| print("run found", run.name) | |
| print(run.summary) | |
| matched_run = run | |
| run_his = matched_run.scan_history() | |
| #print([len(i) for i in run_his]) | |
| for stat in run_his: | |
| stat_i = {k: v for k, v in stat.items() if k not in exclude_keys and v is not None} | |
| if len(stat_i) > 1: | |
| history.append(stat_i) | |
| #matched_run.delete() | |
| break | |
| wandb_run = wandb.init( | |
| project = 'language_det_eval', | |
| job_type = 'evaluate', | |
| name = eval_wandb_name, | |
| ) | |
| #pprint(history) | |
| # exclude_keys = ['_step', '_runtime', '_timestamp'] | |
| # for stat in history: | |
| # wandb.log( | |
| # {k: v for k, v in stat.items() if k not in exclude_keys}, | |
| # step = stat['_step'], | |
| # ) | |
| else: | |
| wandb_run = None | |
| history = None | |
| print("weight_iter: ", weight_iter) | |
| print("train_wandb_name: ", train_wandb_name) | |
| print("eval_wandb_name: ", eval_wandb_name) | |
| if args.task_config: | |
| all_task_configs = args.task_config.split(",") | |
| for task_config in all_task_configs: | |
| cfg_ = cfg.clone() | |
| cfg_.defrost() | |
| cfg_.merge_from_file(task_config) | |
| cfg_.merge_from_list(args.opts) | |
| 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_function( | |
| 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_, | |
| wandb_run=wandb_run, | |
| weight_iter=weight_iter, | |
| history=history, | |
| ) | |
| synchronize() | |
| # logger.info("FLOPs: {}, #Parameter: {}".format(params, flops)) | |
| else: | |
| 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_function( | |
| 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, | |
| wandb_run=wandb_run, | |
| weight_iter=weight_iter, | |
| history=history | |
| ) | |
| synchronize() | |
| # logger.info("FLOPs: {}, #Parameter: {}".format(params, flops)) | |
| if __name__ == "__main__": | |
| main() | |