def set_task(pl_module):
if not pl_module.training:
pl_module.current_tasks = [
k for k, v in pl_module.hparams.config["loss_names"].items()
if v >= 1
]
return
if not check_non_acc_grad(pl_module):
return
sampling_pools = list()
for k, v in pl_module.hparams.config["loss_names"].items():
sampling_pools.extend([k] * v)
current_tasks = [random.choice(sampling_pools)]
picked = all_gather(current_tasks)
pl_module.current_tasks = picked[0]
def compute_irtr_recall(pl_module):
text_dset = pl_module.trainer.datamodule.dms[0].make_no_false_val_dset()
text_dset.tokenizer = pl_module.trainer.datamodule.dms[0].tokenizer
text_loader = torch.utils.data.DataLoader(
text_dset,
batch_size=64,
num_workers=pl_module.hparams.config["num_workers"],
pin_memory=True,
collate_fn=functools.partial(
text_dset.collate,
mlm_collator=pl_module.trainer.datamodule.dms[0].mlm_collator,
),
)
image_dset = pl_module.trainer.datamodule.dms[0].make_no_false_val_dset(
image_only=True)
image_dset.tokenizer = pl_module.trainer.datamodule.dms[0].tokenizer
dist_sampler = DistributedSampler(image_dset, shuffle=False)
image_loader = torch.utils.data.DataLoader(
image_dset,
batch_size=1,
num_workers=pl_module.hparams.config["num_workers"],
sampler=dist_sampler,
pin_memory=True,
collate_fn=functools.partial(
image_dset.collate,
mlm_collator=pl_module.trainer.datamodule.dms[0].mlm_collator,
),
)
text_preload = list()
for _b in tqdm.tqdm(text_loader, desc="text prefetch loop"):
text_preload.append({
"text_ids":
_b["text_ids"].to(pl_module.device),
"text_masks":
_b["text_masks"].to(pl_module.device),
"text_labels":
_b["text_labels"].to(pl_module.device),
"img_index":
_b["img_index"],
})
tiids = list()
for pre in text_preload:
tiids += pre["img_index"]
tiids = torch.tensor(tiids)
image_preload = list()
for _b in tqdm.tqdm(image_loader, desc="image prefetch loop"):
(ie, im, _, _) = pl_module.transformer.visual_embed(
_b["image"][0].to(pl_module.device),
max_image_len=pl_module.hparams.config["max_image_len"],
mask_it=False,
)
image_preload.append((ie, im, _b["img_index"][0]))
rank_scores = list()
rank_iids = list()
for img_batch in tqdm.tqdm(image_preload, desc="rank loop"):
_ie, _im, _iid = img_batch
_, l, c = _ie.shape
img_batch_score = list()
for txt_batch in text_preload:
fblen = len(txt_batch["text_ids"])
ie = _ie.expand(fblen, l, c)
im = _im.expand(fblen, l)
with torch.cuda.amp.autocast():
score = pl_module.rank_output(
pl_module.infer(
{
"text_ids": txt_batch["text_ids"],
"text_masks": txt_batch["text_masks"],
"text_labels": txt_batch["text_labels"],
},
image_embeds=ie,
image_masks=im,
)["cls_feats"])[:, 0]
img_batch_score.append(score)
img_batch_score = torch.cat(img_batch_score)
rank_scores.append(img_batch_score.cpu().tolist())
rank_iids.append(_iid)
torch.distributed.barrier()
gather_rank_scores = all_gather(rank_scores)
gather_rank_iids = all_gather(rank_iids)
iids = torch.tensor(gather_rank_iids)
iids = iids.view(-1)
scores = torch.tensor(gather_rank_scores)
scores = scores.view(len(iids), -1)
topk10 = scores.topk(10, dim=1)
topk5 = scores.topk(5, dim=1)
topk1 = scores.topk(1, dim=1)
topk10_iids = tiids[topk10.indices]
topk5_iids = tiids[topk5.indices]
topk1_iids = tiids[topk1.indices]
tr_r10 = (iids.unsqueeze(1) == topk10_iids).float().max(dim=1)[0].mean()
tr_r5 = (iids.unsqueeze(1) == topk5_iids).float().max(dim=1)[0].mean()
tr_r1 = (iids.unsqueeze(1) == topk1_iids).float().max(dim=1)[0].mean()
topk10 = scores.topk(10, dim=0)
topk5 = scores.topk(5, dim=0)
topk1 = scores.topk(1, dim=0)
topk10_iids = iids[topk10.indices]
topk5_iids = iids[topk5.indices]
topk1_iids = iids[topk1.indices]
ir_r10 = (tiids.unsqueeze(0) == topk10_iids).float().max(dim=0)[0].mean()
ir_r5 = (tiids.unsqueeze(0) == topk5_iids).float().max(dim=0)[0].mean()
ir_r1 = (tiids.unsqueeze(0) == topk1_iids).float().max(dim=0)[0].mean()
return (ir_r1, ir_r5, ir_r10, tr_r1, tr_r5, tr_r10)