| import logging |
| import os |
| from collections import OrderedDict |
| import torch |
| from torch.nn.parallel import DistributedDataParallel |
| import time |
| import datetime |
|
|
| from fvcore.common.timer import Timer |
| import detectron2.utils.comm as comm |
| from detectron2.checkpoint import DetectionCheckpointer, PeriodicCheckpointer |
| from detectron2.config import get_cfg |
| from detectron2.data import ( |
| DatasetCatalog, |
| MetadataCatalog, |
| build_detection_test_loader, |
| ) |
| from detectron2.engine import default_argument_parser, default_setup, launch |
|
|
| from detectron2.evaluation import ( |
| COCOEvaluator, |
| LVISEvaluator, |
| inference_on_dataset, |
| print_csv_format, |
| ) |
| from detectron2.modeling import build_model |
| from detectron2.solver import build_lr_scheduler, build_optimizer |
| from detectron2.utils.events import ( |
| CommonMetricPrinter, |
| EventStorage, |
| JSONWriter, |
| TensorboardXWriter, |
| ) |
| from detectron2.modeling.test_time_augmentation import GeneralizedRCNNWithTTA |
| from detectron2.data.dataset_mapper import DatasetMapper |
| from detectron2.data.build import build_detection_train_loader |
|
|
| from centernet.config import add_centernet_config |
| from centernet.data.custom_build_augmentation import build_custom_augmentation |
| from detectron2.data.datasets import register_coco_instances |
|
|
| logger = logging.getLogger("detectron2") |
|
|
| PATIENCE_EPOCH = 20 |
|
|
| TRAIN_ANN_PATH = 'annotations/instances_train2017.json' |
| TRAIN_IMG_DIR = 'images/instances_train2017/' |
| VAL_ANN_PATH = 'annotations/instances_val2017.json' |
| VAL_IMG_DIR = 'images/instances_val2017/' |
| TEST_ANN_PATH = 'annotations/instances_test2017.json' |
| TEST_IMG_DIR = 'images/instances_test2017/' |
|
|
| def do_test(cfg, model): |
| results = OrderedDict() |
| for dataset_name in cfg.DATASETS.TEST: |
| mapper = None if cfg.INPUT.TEST_INPUT_TYPE == 'default' else \ |
| DatasetMapper( |
| cfg, False, augmentations=build_custom_augmentation(cfg, False)) |
| data_loader = build_detection_test_loader(cfg, dataset_name, mapper=mapper) |
| output_folder = os.path.join( |
| cfg.OUTPUT_DIR, "inference_{}".format(dataset_name)) |
| evaluator_type = MetadataCatalog.get(dataset_name).evaluator_type |
|
|
| if evaluator_type == "lvis": |
| evaluator = LVISEvaluator(dataset_name, cfg, True, output_folder) |
| elif evaluator_type == 'coco': |
| evaluator = COCOEvaluator(dataset_name, cfg, True, output_folder) |
| else: |
| assert 0, evaluator_type |
| |
| results[dataset_name] = inference_on_dataset( |
| model, data_loader, evaluator) |
| if comm.is_main_process(): |
| logger.info("Evaluation results for {} in csv format:".format( |
| dataset_name)) |
| print_csv_format(results[dataset_name]) |
| if len(results) == 1: |
| results = list(results.values())[0] |
| return results |
|
|
|
|
| def do_train(cfg, model, resume=False): |
| model.train() |
| optimizer = build_optimizer(cfg, model) |
| scheduler = build_lr_scheduler(cfg, optimizer) |
|
|
| checkpointer = DetectionCheckpointer( |
| model, cfg.OUTPUT_DIR, optimizer=optimizer, scheduler=scheduler |
| ) |
|
|
| start_iter = ( |
| checkpointer.resume_or_load( |
| cfg.MODEL.WEIGHTS, resume=resume, |
| ).get("iteration", -1) + 1 |
| ) |
| if cfg.SOLVER.RESET_ITER: |
| logger.info('Reset loaded iteration. Start training from iteration 0.') |
| start_iter = 0 |
| max_iter = cfg.SOLVER.MAX_ITER if cfg.SOLVER.TRAIN_ITER < 0 else cfg.SOLVER.TRAIN_ITER |
|
|
| periodic_checkpointer = PeriodicCheckpointer( |
| checkpointer, cfg.SOLVER.CHECKPOINT_PERIOD, max_iter=max_iter |
| ) |
|
|
| writers = ( |
| [ |
| CommonMetricPrinter(max_iter), |
| JSONWriter(os.path.join(cfg.OUTPUT_DIR, "metrics.json")), |
| TensorboardXWriter(cfg.OUTPUT_DIR), |
| ] |
| if comm.is_main_process() |
| else [] |
| ) |
|
|
|
|
| mapper = DatasetMapper(cfg, True) if cfg.INPUT.CUSTOM_AUG == '' else \ |
| DatasetMapper(cfg, True, augmentations=build_custom_augmentation(cfg, True)) |
| if cfg.DATALOADER.SAMPLER_TRAIN in ['TrainingSampler', 'RepeatFactorTrainingSampler']: |
| data_loader = build_detection_train_loader(cfg, mapper=mapper) |
| else: |
| from centernet.data.custom_dataset_dataloader import build_custom_train_loader |
| data_loader = build_custom_train_loader(cfg, mapper=mapper) |
|
|
| iter_per_epoch = len(DatasetCatalog.get("train")) // cfg.SOLVER.IMS_PER_BATCH |
| cfg.defrost() |
| cfg.TEST.EVAL_PERIOD = iter_per_epoch |
| cfg.freeze() |
| patience_iter = PATIENCE_EPOCH * iter_per_epoch |
|
|
| logger.info("Starting training from iteration {}".format(start_iter)) |
| with EventStorage(start_iter) as storage: |
| step_timer = Timer() |
| data_timer = Timer() |
| start_time = time.perf_counter() |
|
|
| best_ap = -1 |
| best_iter = start_iter |
|
|
| def is_nan(x): |
| return x != x |
|
|
| for data, iteration in zip(data_loader, range(start_iter, max_iter)): |
| data_time = data_timer.seconds() |
| storage.put_scalars(data_time=data_time) |
| step_timer.reset() |
| iteration = iteration + 1 |
| storage.step() |
| loss_dict = model(data) |
|
|
| losses = sum( |
| loss for k, loss in loss_dict.items()) |
| assert torch.isfinite(losses).all(), loss_dict |
|
|
| loss_dict_reduced = {k: v.item() \ |
| for k, v in comm.reduce_dict(loss_dict).items()} |
| losses_reduced = sum(loss for loss in loss_dict_reduced.values()) |
| if comm.is_main_process(): |
| storage.put_scalars( |
| total_loss=losses_reduced, **loss_dict_reduced) |
|
|
| optimizer.zero_grad() |
| losses.backward() |
| optimizer.step() |
|
|
| storage.put_scalar( |
| "lr", optimizer.param_groups[0]["lr"], smoothing_hint=False) |
|
|
| step_time = step_timer.seconds() |
| storage.put_scalars(time=step_time) |
| data_timer.reset() |
| scheduler.step() |
|
|
| if ( |
| cfg.TEST.EVAL_PERIOD > 0 |
| and iteration % cfg.TEST.EVAL_PERIOD == 0 |
| and iteration != max_iter |
| ): |
| results = do_test(cfg, model) |
|
|
| if comm.is_main_process(): |
| ap = 0 |
| ap50 = 0 |
| ap75 = 0 |
| if 'bbox' in results: |
| ap = results['bbox']['AP'] |
| ap50 = results['bbox']['AP50'] |
| ap75 = results['bbox']['AP75'] |
| ap = ap if not is_nan(ap) else 0 |
| ap50 = ap50 if not is_nan(ap50) else 0 |
| ap75 = ap75 if not is_nan(ap75) else 0 |
| |
| storage.put_scalar('val/ap', ap) |
| storage.put_scalar('val/ap50', ap50) |
| storage.put_scalar('val/ap75', ap75) |
|
|
| if ap > best_ap: |
| best_ap = ap |
| best_iter = iteration |
| checkpointer.save("model_best") |
| with open(os.path.join(cfg.OUTPUT_DIR, 'best_iter'), 'w') as f: |
| f.write(str(best_iter)) |
| |
| comm.synchronize() |
|
|
| with open(os.path.join(cfg.OUTPUT_DIR, 'best_iter'), 'r') as f: |
| best_iter = int(f.read()) |
|
|
| if iteration - best_iter >= patience_iter: |
| logger.info('No improvement after {} iterations, early stopping at iteration {}, saving best result at iteration {}'.format( |
| patience_iter, |
| iteration, |
| best_iter)) |
| break |
|
|
| if iteration - start_iter > 5 and \ |
| (iteration % 20 == 0 or iteration == max_iter): |
| for writer in writers: |
| writer.write() |
| periodic_checkpointer.step(iteration) |
|
|
| total_time = time.perf_counter() - start_time |
| logger.info( |
| "Total training time: {}".format( |
| str(datetime.timedelta(seconds=int(total_time))))) |
|
|
| def setup(args): |
| """ |
| Create configs and perform basic setups. |
| """ |
| cfg = get_cfg() |
| add_centernet_config(cfg) |
| cfg.merge_from_file(args.config_file) |
| cfg.merge_from_list(args.opts) |
| if '/auto' in cfg.OUTPUT_DIR: |
| file_name = os.path.basename(args.config_file)[:-5] |
| cfg.OUTPUT_DIR = cfg.OUTPUT_DIR.replace('/auto', '/{}'.format(file_name)) |
| logger.info('OUTPUT_DIR: {}'.format(cfg.OUTPUT_DIR)) |
| |
| train_ann_path = os.path.join(args.dataset_root, TRAIN_ANN_PATH) |
| train_img_dir = os.path.join(args.dataset_root, TRAIN_IMG_DIR) |
| val_ann_path = os.path.join(args.dataset_root, VAL_ANN_PATH) |
| val_img_dir = os.path.join(args.dataset_root, VAL_IMG_DIR) |
| test_ann_path = os.path.join(args.dataset_root, TEST_ANN_PATH) |
| test_img_dir = os.path.join(args.dataset_root, TEST_IMG_DIR) |
| register_coco_instances("train", {}, train_ann_path, train_img_dir) |
| register_coco_instances("val", {}, val_ann_path, val_img_dir) |
| register_coco_instances("test", {}, test_ann_path, test_img_dir) |
| cfg.DATASETS.TRAIN = ("train",) |
| |
| cfg.DATASETS.TEST = ("val",) |
| cfg.MODEL.ROI_HEADS.NUM_CLASSES = int(args.num_classes) |
| |
| cfg.freeze() |
| default_setup(cfg, args) |
| return cfg |
|
|
|
|
| def main(args): |
| cfg = setup(args) |
|
|
| model = build_model(cfg) |
| logger.info("Model:\n{}".format(model)) |
|
|
| if args.eval_only: |
| DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load( |
| cfg.MODEL.WEIGHTS, resume=args.resume |
| ) |
| if cfg.TEST.AUG.ENABLED: |
| logger.info("Running inference with test-time augmentation ...") |
| model = GeneralizedRCNNWithTTA(cfg, model, batch_size=1) |
|
|
| return do_test(cfg, model) |
|
|
| distributed = comm.get_world_size() > 1 |
| if distributed: |
| model = DistributedDataParallel( |
| model, device_ids=[comm.get_local_rank()], broadcast_buffers=False, |
| find_unused_parameters=True |
| ) |
|
|
| do_train(cfg, model, resume=args.resume) |
|
|
| model = build_model(cfg) |
| DetectionCheckpointer(model, save_dir=cfg.OUTPUT_DIR).resume_or_load( |
| os.path.join(cfg.OUTPUT_DIR, 'model_best.pth'), resume=False |
| ) |
| cfg.defrost() |
| cfg.DATASETS.TEST = ("test",) |
| cfg.freeze() |
| return do_test(cfg, model) |
|
|
|
|
| if __name__ == "__main__": |
| args = default_argument_parser() |
| args.add_argument('--manual_device', default='') |
| args.add_argument('--dataset_root', default='../../dataset/data/split/interactable/') |
| args.add_argument('--num_classes', default=1) |
| args = args.parse_args() |
| if args.manual_device != '': |
| os.environ['CUDA_VISIBLE_DEVICES'] = args.manual_device |
| args.dist_url = 'tcp://127.0.0.1:{}'.format( |
| torch.randint(11111, 60000, (1,))[0].item()) |
| print("Command Line Args:", args) |
| launch( |
| main, |
| args.num_gpus, |
| num_machines=args.num_machines, |
| machine_rank=args.machine_rank, |
| dist_url=args.dist_url, |
| args=(args,), |
| ) |
|
|
