| import logging | |
| import os | |
| from collections import OrderedDict | |
| import torch | |
| from torch.nn.parallel import DistributedDataParallel | |
| import time | |
| import datetime | |
| import json | |
| 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 ( | |
| 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 | |
| logger = logging.getLogger("detectron2") | |
| 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) | |
| 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() | |
| 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 | |
| ): | |
| do_test(cfg, model) | |
| comm.synchronize() | |
| 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)) | |
| 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) | |
| return do_test(cfg, model) | |
| if __name__ == "__main__": | |
| args = default_argument_parser() | |
| args.add_argument('--manual_device', default='') | |
| 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,), | |
| ) | |