def dump_result(self, path):
if dist.is_main_process():
with open(path, "w") as f:
for p in self.predictions:
f.write(p)
f.write("\n")
def summarize(self):
if dist.is_main_process():
imgid2pred: Dict[str, List] = defaultdict(list)
for p in self.predictions:
imgid2pred[p["original_id"]].append(p)
for img_id, pred in imgid2pred.items():
im_pred_dict = {p["task_id"]: p for p in pred}
self.evaluator.eval_single_img(
img_id,
im_pred_dict,
pred_mask_tag="masks" if self.eval_mask else None,
pred_boxes_tag="boxes",
verbose=False,
)
mask_box = ["box"]
if self.eval_mask:
mask_box.append("mask")
results = self.evaluator.analyze_stats(mask_box,
exp_name_in_summary=None,
save_result_to_path=None)
results = results["all"]["pred_box_acc"]
return {f"Precision@{k}": v for k, v in results.items()}
return None
def summarize(self):
if not dist.is_main_process():
return
if self.fixed_ap:
self._summarize_fixed()
else:
self._summarize_standard()
def __init__(self, topk=10000, fixed_ap=True, out_path="lvis_eval"):
self.results = []
self.by_cat = {}
self.topk = topk
self.fixed_ap = fixed_ap
self.out_path = out_path
if dist.is_main_process():
if not os.path.exists(self.out_path):
os.mkdir(self.out_path)
def __init__(self,
split,
ann_folder,
output_dir="phrasecut_eval",
eval_mask=False):
subset = PhraseCutSubsets(ann_folder)
loader = RefVGLoader(ann_folder, subset, split=split)
if dist.is_main_process():
if not os.path.exists(output_dir):
os.mkdir(output_dir)
self.output_dir = output_dir
self.evaluator = Evaluator(loader, summary_path=output_dir)
self.eval_mask = eval_mask
self.predictions = []
def summarize(self):
if dist.is_main_process():
dataset2score = {
"refcoco": {k: 0.0 for k in self.k},
"refcoco+": {k: 0.0 for k in self.k},
"refcocog": {k: 0.0 for k in self.k},
}
dataset2count = {"refcoco": 0.0, "refcoco+": 0.0, "refcocog": 0.0}
for image_id in self.img_ids:
ann_ids = self.refexp_gt.getAnnIds(imgIds=image_id)
assert len(ann_ids) == 1
img_info = self.refexp_gt.loadImgs(image_id)[0]
target = self.refexp_gt.loadAnns(ann_ids[0])
prediction = self.predictions[image_id]
assert prediction is not None
sorted_scores_boxes = sorted(
zip(prediction["scores"].tolist(), prediction["boxes"].tolist()), reverse=True
)
sorted_scores, sorted_boxes = zip(*sorted_scores_boxes)
sorted_boxes = torch.cat([torch.as_tensor(x).view(1, 4) for x in sorted_boxes])
target_bbox = target[0]["bbox"]
converted_bbox = [
target_bbox[0],
target_bbox[1],
target_bbox[2] + target_bbox[0],
target_bbox[3] + target_bbox[1],
]
giou = generalized_box_iou(sorted_boxes, torch.as_tensor(converted_bbox).view(-1, 4))
for k in self.k:
if max(giou[:k]) >= self.thresh_iou:
dataset2score[img_info["dataset_name"]][k] += 1.0
dataset2count[img_info["dataset_name"]] += 1.0
for key, value in dataset2score.items():
for k in self.k:
try:
value[k] /= dataset2count[key]
except:
pass
results = {}
for key, value in dataset2score.items():
results[key] = sorted([v for k, v in value.items()])
print(f" Dataset: {key} - Precision @ 1, 5, 10: {results[key]} \n")
return results
return None
def summarize(self):
if dist.is_main_process():
self.results = self.evaluator.evaluate(self.predictions)
table = PrettyTable()
all_cat = sorted(list(self.results.values())[0].keys())
table.field_names = ["Recall@k"] + all_cat
score = {}
for k, v in self.results.items():
cur_results = [v[cat] for cat in all_cat]
header = "Upper_bound" if k == -1 else f"Recall@{k}"
for cat in all_cat:
score[f"{header}_{cat}"] = v[cat]
table.add_row([header] + cur_results)
print(table)
return score
return None, None
def dump_result(self, path):
if dist.is_main_process():
with open(path, "w") as f:
json.dump(self.predictions, f, indent=4, sort_keys=True)
def main(args):
# Init distributed mode
dist.init_distributed_mode(args)
# Update dataset specific configs
if args.dataset_config is not None:
# https://stackoverflow.com/a/16878364
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()))
# Segmentation related
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)
# fix the seed for reproducibility
seed = args.seed + dist.get_rank()
torch.manual_seed(seed)
np.random.seed(seed)
random.seed(seed)
torch.set_deterministic(True)
# Build the model
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)"
# Get a copy of the model for exponential moving averaged version of the model
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)
# Set up optimizers
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}")
# Train dataset
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
])
# To handle very big datasets, we chunk it into smaller parts.
if args.epoch_chunks > 0:
print(
"Splitting the training set into {args.epoch_chunks} of size approximately "
f" {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,
)
# Val dataset
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)
# Used for loading weights from another model and starting a training from scratch. Especially useful if
# loading into a model with different functionality.
if args.load:
print("loading from", args.load)
checkpoint = torch.load(args.load, map_location="cpu")
if "model_ema" in checkpoint:
model_without_ddp.load_state_dict(checkpoint["model_ema"],
strict=False)
else:
model_without_ddp.load_state_dict(checkpoint["model"],
strict=False)
if args.ema:
model_ema = deepcopy(model_without_ddp)
# Used for resuming training from the checkpoint of a model. Used when training times-out or is pre-empted.
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
# Runs only evaluation, by default on the validation set unless --test is passed.
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
# Runs training and evaluates after every --eval_skip epochs
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"]
# extra checkpoint before LR drop and every 2 epochs
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:
metric = test_stats["gqa_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"]
# extra checkpoint before LR drop and every 100 epochs
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))
