def main(args):
    dist.init_distributed_mode(args)
    if (args.dataset_config is not None):
        d = vars(args)
        with open(args.dataset_config, 'r') as f:
            cfg = json.load(f)
        d.update(cfg)
    print('git:\n  {}\n'.format(utils.get_sha()))
    if (args.mask_model != 'none'):
        args.masks = True
    if (args.frozen_weights is not None):
        assert args.masks, 'Frozen training is meant for segmentation only'
    print(args)
    device = torch.device(args.device)
    output_dir = Path(args.output_dir)
    seed = (args.seed + dist.get_rank())
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.set_deterministic(True)
    (model, criterion, contrastive_criterion, qa_criterion, weight_dict) = build_model(args)
    model.to(device)
    assert ((criterion is not None) or (qa_criterion is not None)), 'Error: should train either detection or question answering (or both)'
    model_ema = (deepcopy(model) if args.ema else None)
    model_without_ddp = model
    if args.distributed:
        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu], find_unused_parameters=True)
        model_without_ddp = model.module
    n_parameters = sum((p.numel() for p in model.parameters() if p.requires_grad))
    print('number of params:', n_parameters)
    param_dicts = [{'params': [p for (n, p) in model_without_ddp.named_parameters() if (('backbone' not in n) and ('text_encoder' not in n) and p.requires_grad)]}, {'params': [p for (n, p) in model_without_ddp.named_parameters() if (('backbone' in n) and p.requires_grad)], 'lr': args.lr_backbone}, {'params': [p for (n, p) in model_without_ddp.named_parameters() if (('text_encoder' in n) and p.requires_grad)], 'lr': args.text_encoder_lr}]
    if (args.optimizer == 'sgd'):
        optimizer = torch.optim.SGD(param_dicts, lr=args.lr, momentum=0.9, weight_decay=args.weight_decay)
    elif (args.optimizer in ['adam', 'adamw']):
        optimizer = torch.optim.AdamW(param_dicts, lr=args.lr, weight_decay=args.weight_decay)
    else:
        raise RuntimeError(f'Unsupported optimizer {args.optimizer}')
    if ((len(args.combine_datasets) == 0) and (not args.eval)):
        raise RuntimeError('Please provide at least one training dataset')
    (dataset_train, sampler_train, data_loader_train) = (None, None, None)
    if (not args.eval):
        dataset_train = ConcatDataset([build_dataset(name, image_set='train', args=args) for name in args.combine_datasets])
        if (args.epoch_chunks > 0):
            print(f'Splitting the training set into {{args.epoch_chunks}} of size approximately  {(len(dataset_train) // args.epoch_chunks)}')
            chunks = torch.chunk(torch.arange(len(dataset_train)), args.epoch_chunks)
            datasets = [torch.utils.data.Subset(dataset_train, chunk.tolist()) for chunk in chunks]
            if args.distributed:
                samplers_train = [DistributedSampler(ds) for ds in datasets]
            else:
                samplers_train = [torch.utils.data.RandomSampler(ds) for ds in datasets]
            batch_samplers_train = [torch.utils.data.BatchSampler(sampler_train, args.batch_size, drop_last=True) for sampler_train in samplers_train]
            assert (len(batch_samplers_train) == len(datasets))
            data_loaders_train = [DataLoader(ds, batch_sampler=batch_sampler_train, collate_fn=partial(utils.collate_fn, False), num_workers=args.num_workers) for (ds, batch_sampler_train) in zip(datasets, batch_samplers_train)]
        else:
            if args.distributed:
                sampler_train = DistributedSampler(dataset_train)
            else:
                sampler_train = torch.utils.data.RandomSampler(dataset_train)
            batch_sampler_train = torch.utils.data.BatchSampler(sampler_train, args.batch_size, drop_last=True)
            data_loader_train = DataLoader(dataset_train, batch_sampler=batch_sampler_train, collate_fn=partial(utils.collate_fn, False), num_workers=args.num_workers)
    if (len(args.combine_datasets_val) == 0):
        raise RuntimeError('Please provide at leas one validation dataset')
    Val_all = namedtuple(typename='val_data', field_names=['dataset_name', 'dataloader', 'base_ds', 'evaluator_list'])
    val_tuples = []
    for dset_name in args.combine_datasets_val:
        dset = build_dataset(dset_name, image_set='val', args=args)
        sampler = (DistributedSampler(dset, shuffle=False) if args.distributed else torch.utils.data.SequentialSampler(dset))
        dataloader = DataLoader(dset, args.batch_size, sampler=sampler, drop_last=False, collate_fn=partial(utils.collate_fn, False), num_workers=args.num_workers)
        base_ds = get_coco_api_from_dataset(dset)
        val_tuples.append(Val_all(dataset_name=dset_name, dataloader=dataloader, base_ds=base_ds, evaluator_list=None))
    if (args.frozen_weights is not None):
        if args.resume.startswith('https'):
            checkpoint = torch.hub.load_state_dict_from_url(args.resume, map_location='cpu', check_hash=True)
        else:
            checkpoint = torch.load(args.resume, map_location='cpu')
        if (('model_ema' in checkpoint) and (checkpoint['model_ema'] is not None)):
            model_without_ddp.detr.load_state_dict(checkpoint['model_ema'], strict=False)
        else:
            model_without_ddp.detr.load_state_dict(checkpoint['model'], strict=False)
        if args.ema:
            model_ema = deepcopy(model_without_ddp)
    if args.load:
        print('loading from', args.load)
        checkpoint = torch.load(args.load, map_location='cpu')
        if ('model_ema' in checkpoint):
            if args.do_qa_with_qa_fine_tuned:
                del checkpoint['model_ema']['qa_embed.weight']
                del checkpoint['model_ema']['answer_type_head.weight']
                del checkpoint['model_ema']['answer_type_head.bias']
            model_without_ddp.load_state_dict(checkpoint['model_ema'], strict=False)
        else:
            if args.do_qa_with_qa_fine_tuned:
                del checkpoint['model']['qa_embed.weight']
                del checkpoint['model']['answer_type_head.weight']
                del checkpoint['model']['answer_type_head.bias']
            model_without_ddp.load_state_dict(checkpoint['model'], strict=False)
        if args.ema:
            model_ema = deepcopy(model_without_ddp)
    if args.resume:
        if args.resume.startswith('https'):
            checkpoint = torch.hub.load_state_dict_from_url(args.resume, map_location='cpu', check_hash=True)
        else:
            checkpoint = torch.load(args.resume, map_location='cpu')
        model_without_ddp.load_state_dict(checkpoint['model'])
        if ((not args.eval) and ('optimizer' in checkpoint) and ('epoch' in checkpoint)):
            optimizer.load_state_dict(checkpoint['optimizer'])
            args.start_epoch = (checkpoint['epoch'] + 1)
        if args.ema:
            if ('model_ema' not in checkpoint):
                print('WARNING: ema model not found in checkpoint, resetting to current model')
                model_ema = deepcopy(model_without_ddp)
            else:
                model_ema.load_state_dict(checkpoint['model_ema'])

    def build_evaluator_list(base_ds, dataset_name):
        'Helper function to build the list of evaluators for a given dataset'
        evaluator_list = []
        if args.no_detection:
            return evaluator_list
        iou_types = ['bbox']
        if args.masks:
            iou_types.append('segm')
        evaluator_list.append(CocoEvaluator(base_ds, tuple(iou_types), useCats=False))
        if ('refexp' in dataset_name):
            evaluator_list.append(RefExpEvaluator(base_ds, 'bbox'))
        if ('clevrref' in dataset_name):
            evaluator_list.append(ClevrRefEvaluator(base_ds, 'bbox'))
        if ('flickr' in dataset_name):
            evaluator_list.append(FlickrEvaluator(args.flickr_dataset_path, subset=('test' if args.test else 'val'), merge_boxes=(args.GT_type == 'merged')))
        if ('phrasecut' in dataset_name):
            evaluator_list.append(PhrasecutEvaluator(('test' if args.test else 'miniv'), ann_folder=args.phrasecut_orig_ann_path, output_dir=os.path.join(output_dir, 'phrasecut_eval'), eval_mask=args.masks))
        return evaluator_list
    if args.eval:
        test_stats = {}
        test_model = (model_ema if (model_ema is not None) else model)
        for (i, item) in enumerate(val_tuples):
            evaluator_list = build_evaluator_list(item.base_ds, item.dataset_name)
            postprocessors = build_postprocessors(args, item.dataset_name)
            item = item._replace(evaluator_list=evaluator_list)
            print(f'Evaluating {item.dataset_name}')
            curr_test_stats = evaluate(model=test_model, criterion=criterion, contrastive_criterion=contrastive_criterion, qa_criterion=qa_criterion, postprocessors=postprocessors, weight_dict=weight_dict, data_loader=item.dataloader, evaluator_list=item.evaluator_list, device=device, args=args)
            test_stats.update({((item.dataset_name + '_') + k): v for (k, v) in curr_test_stats.items()})
        log_stats = {**{f'test_{k}': v for (k, v) in test_stats.items()}, 'n_parameters': n_parameters}
        print(log_stats)
        return
    print('Start training')
    start_time = time.time()
    best_metric = 0.0
    for epoch in range(args.start_epoch, args.epochs):
        if (args.epoch_chunks > 0):
            sampler_train = samplers_train[(epoch % len(samplers_train))]
            data_loader_train = data_loaders_train[(epoch % len(data_loaders_train))]
            print(f'Starting epoch {(epoch // len(data_loaders_train))}, sub_epoch {(epoch % len(data_loaders_train))}')
        else:
            print(f'Starting epoch {epoch}')
        if args.distributed:
            sampler_train.set_epoch(epoch)
        train_stats = train_one_epoch(model=model, criterion=criterion, contrastive_criterion=contrastive_criterion, qa_criterion=qa_criterion, data_loader=data_loader_train, weight_dict=weight_dict, optimizer=optimizer, device=device, epoch=epoch, args=args, max_norm=args.clip_max_norm, model_ema=model_ema)
        if args.output_dir:
            checkpoint_paths = [(output_dir / 'checkpoint.pth')]
            if ((((epoch + 1) % args.lr_drop) == 0) or (((epoch + 1) % 2) == 0)):
                checkpoint_paths.append((output_dir / f'checkpoint{epoch:04}.pth'))
            for checkpoint_path in checkpoint_paths:
                dist.save_on_master({'model': model_without_ddp.state_dict(), 'model_ema': (model_ema.state_dict() if args.ema else None), 'optimizer': optimizer.state_dict(), 'epoch': epoch, 'args': args}, checkpoint_path)
        if ((epoch % args.eval_skip) == 0):
            test_stats = {}
            test_model = (model_ema if (model_ema is not None) else model)
            for (i, item) in enumerate(val_tuples):
                evaluator_list = build_evaluator_list(item.base_ds, item.dataset_name)
                item = item._replace(evaluator_list=evaluator_list)
                postprocessors = build_postprocessors(args, item.dataset_name)
                print(f'Evaluating {item.dataset_name}')
                curr_test_stats = evaluate(model=test_model, criterion=criterion, contrastive_criterion=contrastive_criterion, qa_criterion=qa_criterion, postprocessors=postprocessors, weight_dict=weight_dict, data_loader=item.dataloader, evaluator_list=item.evaluator_list, device=device, args=args)
                test_stats.update({((item.dataset_name + '_') + k): v for (k, v) in curr_test_stats.items()})
        else:
            test_stats = {}
        log_stats = {**{f'train_{k}': v for (k, v) in train_stats.items()}, **{f'test_{k}': v for (k, v) in test_stats.items()}, 'epoch': epoch, 'n_parameters': n_parameters}
        if (args.output_dir and dist.is_main_process()):
            with (output_dir / 'log.txt').open('a') as f:
                f.write((json.dumps(log_stats) + '\n'))
        if ((epoch % args.eval_skip) == 0):
            if args.do_qa:
                if ('gqa' in args.combine_datasets):
                    metric = test_stats['gqa_accuracy_answer_total_unscaled']
                else:
                    metric = test_stats['vqa2_accuracy_answer_total_unscaled']
            else:
                metric = np.mean([v[1] for (k, v) in test_stats.items() if ('coco_eval_bbox' in k)])
            if (args.output_dir and (metric > best_metric)):
                best_metric = metric
                checkpoint_paths = [(output_dir / 'BEST_checkpoint.pth')]
                for checkpoint_path in checkpoint_paths:
                    dist.save_on_master({'model': model_without_ddp.state_dict(), 'optimizer': optimizer.state_dict(), 'epoch': epoch, 'args': args}, checkpoint_path)
    total_time = (time.time() - start_time)
    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
    print('Training time {}'.format(total_time_str))
def main(args):
    dist.init_distributed_mode(args)
    if (args.dataset_config is not None):
        d = vars(args)
        with open(args.dataset_config, 'r') as f:
            cfg = json.load(f)
        d.update(cfg)
    print('git:\n  {}\n'.format(utils.get_sha()))
    if (args.mask_model != 'none'):
        args.masks = True
    if (args.frozen_weights is not None):
        assert args.masks, 'Frozen training is meant for segmentation only'
    print(args)
    device = torch.device(args.device)
    output_dir = Path(args.output_dir)
    seed = (args.seed + dist.get_rank())
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.set_deterministic(True)
    (model, criterion, contrastive_criterion, qa_criterion, weight_dict) = build_model(args)
    model.to(device)
    assert ((criterion is not None) or (qa_criterion is not None)), 'Error: should train either detection or question answering (or both)'
    model_ema = (deepcopy(model) if args.ema else None)
    model_without_ddp = model
    if args.distributed:
        model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu], find_unused_parameters=True)
        model_without_ddp = model.module
    n_parameters = sum((p.numel() for p in model.parameters() if p.requires_grad))
    print('number of params:', n_parameters)
    param_dicts = [{'params': [p for (n, p) in model_without_ddp.named_parameters() if (('backbone' not in n) and ('text_encoder' not in n) and p.requires_grad)]}, {'params': [p for (n, p) in model_without_ddp.named_parameters() if (('backbone' in n) and p.requires_grad)], 'lr': args.lr_backbone}, {'params': [p for (n, p) in model_without_ddp.named_parameters() if (('text_encoder' in n) and p.requires_grad)], 'lr': args.text_encoder_lr}]
    if (args.optimizer == 'sgd'):
        optimizer = torch.optim.SGD(param_dicts, lr=args.lr, momentum=0.9, weight_decay=args.weight_decay)
    elif (args.optimizer in ['adam', 'adamw']):
        optimizer = torch.optim.AdamW(param_dicts, lr=args.lr, weight_decay=args.weight_decay)
    else:
        raise RuntimeError(f'Unsupported optimizer {args.optimizer}')
    if ((len(args.combine_datasets) == 0) and (not args.eval)):
        raise RuntimeError('Please provide at least one training dataset')
    (dataset_train, sampler_train, data_loader_train) = (None, None, None)
    if (not args.eval):
        dataset_train = ConcatDataset([build_dataset(name, image_set='train', args=args) for name in args.combine_datasets])
        if (args.epoch_chunks > 0):
            print(f'Splitting the training set into {{args.epoch_chunks}} of size approximately  {(len(dataset_train) // args.epoch_chunks)}')
            chunks = torch.chunk(torch.arange(len(dataset_train)), args.epoch_chunks)
            datasets = [torch.utils.data.Subset(dataset_train, chunk.tolist()) for chunk in chunks]
            if args.distributed:
                samplers_train = [DistributedSampler(ds) for ds in datasets]
            else:
                samplers_train = [torch.utils.data.RandomSampler(ds) for ds in datasets]
            batch_samplers_train = [torch.utils.data.BatchSampler(sampler_train, args.batch_size, drop_last=True) for sampler_train in samplers_train]
            assert (len(batch_samplers_train) == len(datasets))
            data_loaders_train = [DataLoader(ds, batch_sampler=batch_sampler_train, collate_fn=partial(utils.collate_fn, False), num_workers=args.num_workers) for (ds, batch_sampler_train) in zip(datasets, batch_samplers_train)]
        else:
            if args.distributed:
                sampler_train = DistributedSampler(dataset_train)
            else:
                sampler_train = torch.utils.data.RandomSampler(dataset_train)
            batch_sampler_train = torch.utils.data.BatchSampler(sampler_train, args.batch_size, drop_last=True)
            data_loader_train = DataLoader(dataset_train, batch_sampler=batch_sampler_train, collate_fn=partial(utils.collate_fn, False), num_workers=args.num_workers)
    if (len(args.combine_datasets_val) == 0):
        raise RuntimeError('Please provide at leas one validation dataset')
    Val_all = namedtuple(typename='val_data', field_names=['dataset_name', 'dataloader', 'base_ds', 'evaluator_list'])
    val_tuples = []
    for dset_name in args.combine_datasets_val:
        dset = build_dataset(dset_name, image_set='val', args=args)
        sampler = (DistributedSampler(dset, shuffle=False) if args.distributed else torch.utils.data.SequentialSampler(dset))
        dataloader = DataLoader(dset, args.batch_size, sampler=sampler, drop_last=False, collate_fn=partial(utils.collate_fn, False), num_workers=args.num_workers)
        base_ds = get_coco_api_from_dataset(dset)
        val_tuples.append(Val_all(dataset_name=dset_name, dataloader=dataloader, base_ds=base_ds, evaluator_list=None))
    if (args.frozen_weights is not None):
        if args.resume.startswith('https'):
            checkpoint = torch.hub.load_state_dict_from_url(args.resume, map_location='cpu', check_hash=True)
        else:
            checkpoint = torch.load(args.resume, map_location='cpu')
        if (('model_ema' in checkpoint) and (checkpoint['model_ema'] is not None)):
            model_without_ddp.detr.load_state_dict(checkpoint['model_ema'], strict=False)
        else:
            model_without_ddp.detr.load_state_dict(checkpoint['model'], strict=False)
        if args.ema:
            model_ema = deepcopy(model_without_ddp)
    if args.load:
        print('loading from', args.load)
        checkpoint = torch.load(args.load, map_location='cpu')
        if ('model_ema' in checkpoint):
            if args.do_qa_with_qa_fine_tuned:
                del checkpoint['model_ema']['qa_embed.weight']
                del checkpoint['model_ema']['answer_type_head.weight']
                del checkpoint['model_ema']['answer_type_head.bias']
            model_without_ddp.load_state_dict(checkpoint['model_ema'], strict=False)
        else:
            if args.do_qa_with_qa_fine_tuned:
                del checkpoint['model']['qa_embed.weight']
                del checkpoint['model']['answer_type_head.weight']
                del checkpoint['model']['answer_type_head.bias']
            model_without_ddp.load_state_dict(checkpoint['model'], strict=False)
        if args.ema:
            model_ema = deepcopy(model_without_ddp)
    if args.resume:
        if args.resume.startswith('https'):
            checkpoint = torch.hub.load_state_dict_from_url(args.resume, map_location='cpu', check_hash=True)
        else:
            checkpoint = torch.load(args.resume, map_location='cpu')
        model_without_ddp.load_state_dict(checkpoint['model'])
        if ((not args.eval) and ('optimizer' in checkpoint) and ('epoch' in checkpoint)):
            optimizer.load_state_dict(checkpoint['optimizer'])
            args.start_epoch = (checkpoint['epoch'] + 1)
        if args.ema:
            if ('model_ema' not in checkpoint):
                print('WARNING: ema model not found in checkpoint, resetting to current model')
                model_ema = deepcopy(model_without_ddp)
            else:
                model_ema.load_state_dict(checkpoint['model_ema'])

    def build_evaluator_list(base_ds, dataset_name):
        'Helper function to build the list of evaluators for a given dataset'
        evaluator_list = []
        if args.no_detection:
            return evaluator_list
        iou_types = ['bbox']
        if args.masks:
            iou_types.append('segm')
        evaluator_list.append(CocoEvaluator(base_ds, tuple(iou_types), useCats=False))
        if ('refexp' in dataset_name):
            evaluator_list.append(RefExpEvaluator(base_ds, 'bbox'))
        if ('clevrref' in dataset_name):
            evaluator_list.append(ClevrRefEvaluator(base_ds, 'bbox'))
        if ('flickr' in dataset_name):
            evaluator_list.append(FlickrEvaluator(args.flickr_dataset_path, subset=('test' if args.test else 'val'), merge_boxes=(args.GT_type == 'merged')))
        if ('phrasecut' in dataset_name):
            evaluator_list.append(PhrasecutEvaluator(('test' if args.test else 'miniv'), ann_folder=args.phrasecut_orig_ann_path, output_dir=os.path.join(output_dir, 'phrasecut_eval'), eval_mask=args.masks))
        return evaluator_list
    if args.eval:
        test_stats = {}
        test_model = (model_ema if (model_ema is not None) else model)
        for (i, item) in enumerate(val_tuples):
            evaluator_list = build_evaluator_list(item.base_ds, item.dataset_name)
            postprocessors = build_postprocessors(args, item.dataset_name)
            item = item._replace(evaluator_list=evaluator_list)
            print(f'Evaluating {item.dataset_name}')
            curr_test_stats = evaluate(model=test_model, criterion=criterion, contrastive_criterion=contrastive_criterion, qa_criterion=qa_criterion, postprocessors=postprocessors, weight_dict=weight_dict, data_loader=item.dataloader, evaluator_list=item.evaluator_list, device=device, args=args)
            test_stats.update({((item.dataset_name + '_') + k): v for (k, v) in curr_test_stats.items()})
        log_stats = {**{f'test_{k}': v for (k, v) in test_stats.items()}, 'n_parameters': n_parameters}
        print(log_stats)
        return
    print('Start training')
    start_time = time.time()
    best_metric = 0.0
    for epoch in range(args.start_epoch, args.epochs):
        if (args.epoch_chunks > 0):
            sampler_train = samplers_train[(epoch % len(samplers_train))]
            data_loader_train = data_loaders_train[(epoch % len(data_loaders_train))]
            print(f'Starting epoch {(epoch // len(data_loaders_train))}, sub_epoch {(epoch % len(data_loaders_train))}')
        else:
            print(f'Starting epoch {epoch}')
        if args.distributed:
            sampler_train.set_epoch(epoch)
        train_stats = train_one_epoch(model=model, criterion=criterion, contrastive_criterion=contrastive_criterion, qa_criterion=qa_criterion, data_loader=data_loader_train, weight_dict=weight_dict, optimizer=optimizer, device=device, epoch=epoch, args=args, max_norm=args.clip_max_norm, model_ema=model_ema)
        if args.output_dir:
            checkpoint_paths = [(output_dir / 'checkpoint.pth')]
            if ((((epoch + 1) % args.lr_drop) == 0) or (((epoch + 1) % 2) == 0)):
                checkpoint_paths.append((output_dir / f'checkpoint{epoch:04}.pth'))
            for checkpoint_path in checkpoint_paths:
                dist.save_on_master({'model': model_without_ddp.state_dict(), 'model_ema': (model_ema.state_dict() if args.ema else None), 'optimizer': optimizer.state_dict(), 'epoch': epoch, 'args': args}, checkpoint_path)
        if ((epoch % args.eval_skip) == 0):
            test_stats = {}
            test_model = (model_ema if (model_ema is not None) else model)
            for (i, item) in enumerate(val_tuples):
                evaluator_list = build_evaluator_list(item.base_ds, item.dataset_name)
                item = item._replace(evaluator_list=evaluator_list)
                postprocessors = build_postprocessors(args, item.dataset_name)
                print(f'Evaluating {item.dataset_name}')
                curr_test_stats = evaluate(model=test_model, criterion=criterion, contrastive_criterion=contrastive_criterion, qa_criterion=qa_criterion, postprocessors=postprocessors, weight_dict=weight_dict, data_loader=item.dataloader, evaluator_list=item.evaluator_list, device=device, args=args)
                test_stats.update({((item.dataset_name + '_') + k): v for (k, v) in curr_test_stats.items()})
        else:
            test_stats = {}
        log_stats = {**{f'train_{k}': v for (k, v) in train_stats.items()}, **{f'test_{k}': v for (k, v) in test_stats.items()}, 'epoch': epoch, 'n_parameters': n_parameters}
        if (args.output_dir and dist.is_main_process()):
            with (output_dir / 'log.txt').open('a') as f:
                f.write((json.dumps(log_stats) + '\n'))
        if ((epoch % args.eval_skip) == 0):
            if args.do_qa:
                if ('gqa' in args.combine_datasets):
                    metric = test_stats['gqa_accuracy_answer_total_unscaled']
                else:
                    metric = test_stats['vqa2_accuracy_answer_total_unscaled']
            else:
                metric = np.mean([v[1] for (k, v) in test_stats.items() if ('coco_eval_bbox' in k)])
            if (args.output_dir and (metric > best_metric)):
                best_metric = metric
                checkpoint_paths = [(output_dir / 'BEST_checkpoint.pth')]
                for checkpoint_path in checkpoint_paths:
                    dist.save_on_master({'model': model_without_ddp.state_dict(), 'optimizer': optimizer.state_dict(), 'epoch': epoch, 'args': args}, checkpoint_path)
    total_time = (time.time() - start_time)
    total_time_str = str(datetime.timedelta(seconds=int(total_time)))
    print('Training time {}'.format(total_time_str))