def train(cfg): model = build_segmentation_model(cfg) device = torch.device(cfg.MODEL.DEVICE) model.to(device) optimizer = make_optimizer(cfg, model) # scheduler = make_lr_scheduler(cfg, optimizer) scheduler = None arguments = {} arguments["epoch"] = 0 output_dir = cfg.OUTPUT_DIR max_epoch = cfg.SOLVER.MAX_EPOCH checkpointer = SegmentationCheckpointer( cfg, model, optimizer, scheduler, output_dir, save_to_disk=True ) extra_checkpoint_data = checkpointer.load(cfg.MODEL.WEIGHT) arguments.update(extra_checkpoint_data) train_data_loader = make_data_loader( cfg, split='train' ) val_data_loader = make_data_loader( cfg, split='val' ) checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD do_train( cfg, model, train_data_loader, val_data_loader, optimizer, scheduler, checkpointer, device, checkpoint_period, arguments, max_epoch, ) return model
def train(cfg, local_rank, distributed): model = build_detection_model(cfg) device = torch.device(cfg.MODEL.DEVICE) model.to(device) if cfg.MODEL.USE_SYNCBN: assert is_pytorch_1_1_0_or_later(), \ "SyncBatchNorm is only available in pytorch >= 1.1.0" model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model) optimizer = make_optimizer(cfg, model) scheduler = make_lr_scheduler(cfg, optimizer) if distributed: model = torch.nn.parallel.DistributedDataParallel( model, device_ids=[local_rank], output_device=local_rank, # this should be removed if we update BatchNorm stats broadcast_buffers=False, # find_unused_parameters=True, ) arguments = {} arguments["iteration"] = 0 output_dir = cfg.OUTPUT_DIR save_to_disk = get_rank() == 0 checkpointer = DetectronCheckpointer(cfg, model, optimizer, scheduler, output_dir, save_to_disk) extra_checkpoint_data = checkpointer.load(cfg.MODEL.WEIGHT) arguments.update(extra_checkpoint_data) data_loader = make_data_loader( cfg, is_train=True, is_distributed=distributed, start_iter=arguments["iteration"], ) checkpoint_period = cfg.SOLVER.CHECKPOINT_PERIOD vis_period = cfg.VISUALIZE.PERIOD if 0 < vis_period < cfg.SOLVER.MAX_ITER: visualizer = SummaryWriterX( cfg.VISUALIZE.DIR + '/' + cfg.VISUALIZE.ENV, cfg.VISUALIZE.ENV, vis_period, 20, get_category(cfg.DATASETS.TRAIN[0])) else: visualizer = None meters = MetricLogger(delimiter=" ", save_dir=os.path.join(output_dir, 'meters.json')) meters.load(is_main_process=get_rank() == 0) do_train(model, data_loader, optimizer, scheduler, checkpointer, device, checkpoint_period, arguments, meters, visualizer) return model
def main(): parser = argparse.ArgumentParser(description="PyTorch Segmentation Inference") parser.add_argument( "--config-file", default="./configs/Encoder_UNet.yaml", metavar="FILE", help="path to config file", ) parser.add_argument( "opts", help="Modify config options using the command-line", default=None, nargs=argparse.REMAINDER, ) args = parser.parse_args() cfg.merge_from_file(args.config_file) cfg.merge_from_list(args.opts) cfg.freeze() save_dir = cfg.OUTPUT_DIR logger = setup_logger("core", save_dir) logger.info(cfg) logger.info("Collecting env info (might take some time)") logger.info("\n" + collect_env_info()) model = build_segmentation_model(cfg) model.to(cfg.MODEL.DEVICE) output_dir = cfg.OUTPUT_DIR checkpointer = SegmentationCheckpointer(cfg, model, save_dir=output_dir) _ = checkpointer.load(cfg.MODEL.WEIGHT) dataset_names = cfg.DATASETS.TEST output_folders = [None] * len(cfg.DATASETS.TEST) if cfg.OUTPUT_DIR: for idx, dataset_name in enumerate(dataset_names): output_folder = os.path.join(cfg.OUTPUT_DIR, "inference", cfg.MODEL.ENCODER + '_' + cfg.MODEL.ARCHITECTURE, dataset_name) mkdir(output_folder) output_folders[idx] = output_folder else: raise RuntimeError("Output directory is missing!") test_data_loaders = make_data_loader(cfg, split='test') for output_folder, dataset_name, test_data_loader in zip(output_folders, dataset_names, test_data_loaders): inference( model, test_data_loader, dataset_name=dataset_name, device=cfg.MODEL.DEVICE, output_folder=output_folder, )
def run_test(cfg, model): dataset_names = cfg.DATASETS.TEST if cfg.OUTPUT_DIR: for idx, dataset_name in enumerate(dataset_names): output_folder = os.path.join(cfg.OUTPUT_DIR, "inference", dataset_name) mkdir(output_folder) output_folders[idx] = output_folder else: raise RuntimeError("Output directory is missing!") test_data_loaders = make_data_loader(cfg, split='test') for output_folder, dataset_name, test_data_loader in zip(output_folders, dataset_names, test_data_loaders): inference( model, test_data_loader, dataset_name=dataset_name, device=cfg.MODEL.DEVICE, output_folder=output_folder, )
def run_test(cfg, model, distributed): if distributed: 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", ) output_folders = [None] * len(cfg.DATASETS.TEST) dataset_names = cfg.DATASETS.TEST if cfg.OUTPUT_DIR: for idx, dataset_name in enumerate(dataset_names): output_folder = os.path.join(cfg.OUTPUT_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=False if cfg.MODEL.FCOS_ON or cfg.MODEL.PACKDET_ON or cfg.MODEL.RETINAPACK_ON or cfg.MODEL.RETINANET_ON else cfg.MODEL.RPN_ONLY, device=cfg.MODEL.DEVICE, expected_results=cfg.TEST.EXPECTED_RESULTS, expected_results_sigma_tol=cfg.TEST.EXPECTED_RESULTS_SIGMA_TOL, output_folder=output_folder, ) 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("--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://" ) synchronize() cfg.merge_from_file(args.config_file) cfg.merge_from_list(args.opts) cfg.freeze() save_dir = "" logger = setup_logger("core", save_dir, get_rank()) 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) num_parameters = sum([param.nelement() for param in model.parameters()]) logger.info('# parameters totally: '+str(num_parameters)) output_dir = cfg.OUTPUT_DIR checkpointer = DetectronCheckpointer(cfg, model, save_dir=output_dir) _ = checkpointer.load(cfg.MODEL.WEIGHT, is_train=False) suffix = cfg.MODEL.WEIGHT.split('/')[-1][:-4] iou_types = ("bbox",) if cfg.MODEL.MASK_ON: iou_types = iou_types + ("segm",) if cfg.MODEL.KEYPOINT_ON: iou_types = iou_types + ("keypoints",) output_folders = [None] * len(cfg.DATASETS.TEST) dataset_names = cfg.DATASETS.TEST if cfg.OUTPUT_DIR: for idx, dataset_name in enumerate(dataset_names): output_folder = os.path.join(cfg.OUTPUT_DIR, "inference_"+suffix, 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=False if cfg.MODEL.FCOS_ON or cfg.MODEL.PACKDET_ON or cfg.MODEL.RETINAPACK_ON or cfg.MODEL.RETINANET_ON else cfg.MODEL.RPN_ONLY, device=cfg.MODEL.DEVICE, expected_results=cfg.TEST.EXPECTED_RESULTS, expected_results_sigma_tol=cfg.TEST.EXPECTED_RESULTS_SIGMA_TOL, output_folder=output_folder, ) synchronize()