add vlmo

vmlo/vlmo/datamodules/vg_caption_datamodule.py

+from vlmo.datasets import VisualGenomeCaptionDataset
+from .datamodule_base import BaseDataModule
+
+
+class VisualGenomeCaptionDataModule(BaseDataModule):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    @property
+    def dataset_cls(self):
+        return VisualGenomeCaptionDataset
+
+    @property
+    def dataset_name(self):
+        return "vg"

vmlo/vlmo/datamodules/vqav2_datamodule.py

+from vlmo.datasets import VQAv2Dataset
+from .datamodule_base import BaseDataModule
+from collections import defaultdict
+
+
+class VQAv2DataModule(BaseDataModule):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    @property
+    def dataset_cls(self):
+        return VQAv2Dataset
+
+    @property
+    def dataset_name(self):
+        return "vqa"
+
+    def setup(self, stage):
+        super().setup(stage)
+
+        train_answers = self.train_dataset.table["answers"].to_pandas().tolist()
+        val_answers = self.val_dataset.table["answers"].to_pandas().tolist()
+        train_labels = self.train_dataset.table["answer_labels"].to_pandas().tolist()
+        val_labels = self.val_dataset.table["answer_labels"].to_pandas().tolist()
+
+        all_answers = [c for c in train_answers + val_answers if c is not None]
+        all_answers = [l for lll in all_answers for ll in lll for l in ll]
+        all_labels = [c for c in train_labels + val_labels if c is not None]
+        all_labels = [l for lll in all_labels for ll in lll for l in ll]
+
+        self.answer2id = {k: v for k, v in zip(all_answers, all_labels)}
+        sorted_a2i = sorted(self.answer2id.items(), key=lambda x: x[1])
+        self.num_class = max(self.answer2id.values()) + 1
+        self.id2answer = defaultdict(lambda: "unknown")
+        for k, v in sorted_a2i:
+            self.id2answer[v] = k

vmlo/vlmo/datamodules/wikibk_datamodule.py

+from vlmo.datasets import WikibkDataset
+from .datamodule_base import BaseDataModule
+
+
+class WikibkDataModule(BaseDataModule):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    @property
+    def dataset_cls(self):
+        return WikibkDataset
+
+    @property
+    def dataset_name(self):
+        return "wikibk"

vmlo/vlmo/datasets/__init__.py

+from .vg_caption_dataset import VisualGenomeCaptionDataset
+from .coco_caption_karpathy_dataset import CocoCaptionKarpathyDataset
+from .f30k_caption_karpathy_dataset import F30KCaptionKarpathyDataset
+from .conceptual_caption_dataset import ConceptualCaptionDataset
+from .sbu_caption_dataset import SBUCaptionDataset
+from .wikibk_dataset import WikibkDataset
+from .vqav2_dataset import VQAv2Dataset
+from .nlvr2_dataset import NLVR2Dataset

vmlo/vlmo/datasets/base_dataset.py

+import random
+import torch
+import io
+import pyarrow as pa
+import os
+
+from PIL import Image
+from vlmo.transforms import keys_to_transforms
+
+
+class BaseDataset(torch.utils.data.Dataset):
+    def __init__(
+        self,
+        data_dir: str,
+        transform_keys: list,
+        image_size: int,
+        names: list,
+        text_column_name: str = "",
+        remove_duplicate=False,
+        max_text_len=40,
+        draw_false_image=0,
+        draw_false_text=0,
+        image_only=False,
+    ):
+        """
+        data_dir : where dataset file *.arrow lives; existence should be guaranteed via DataModule.prepare_data
+        transform_keys : keys for generating augmented views of images
+        text_column_name : pyarrow table column name that has list of strings as elements
+        """
+        assert len(transform_keys) >= 1
+        super().__init__()
+
+        self.transforms = keys_to_transforms(transform_keys, size=image_size)
+        self.text_column_name = text_column_name
+        self.names = names
+        self.max_text_len = max_text_len
+        self.draw_false_image = draw_false_image
+        self.draw_false_text = draw_false_text
+        self.image_only = image_only
+        self.data_dir = data_dir
+
+        if len(names) != 0:
+            tables = [
+                pa.ipc.RecordBatchFileReader(
+                    pa.memory_map(f"{data_dir}/{name}.arrow", "r")
+                ).read_all()
+                for name in names
+                if os.path.isfile(f"{data_dir}/{name}.arrow")
+            ]
+
+            self.table_names = list()
+            for i, name in enumerate(names):
+                self.table_names += [name] * len(tables[i])
+
+            self.table = pa.concat_tables(tables, promote=True)
+            if text_column_name != "":
+                self.text_column_name = text_column_name
+                self.all_texts = self.table[text_column_name].to_pandas().tolist()
+                self.all_texts = (
+                    [list(set(texts)) for texts in self.all_texts]
+                    if remove_duplicate
+                    else self.all_texts
+                )
+            else:
+                self.all_texts = list()
+        else:
+            self.all_texts = list()
+
+        self.index_mapper = dict()
+
+        if text_column_name != "" and not self.image_only:
+            j = 0
+            for i, texts in enumerate(self.all_texts):
+                for _j in range(len(texts)):
+                    self.index_mapper[j] = (i, _j)
+                    j += 1
+        else:
+            for i in range(len(self.table)):
+                self.index_mapper[i] = (i, None)
+
+    @property
+    def corpus(self):
+        return [text for texts in self.all_texts for text in texts]
+
+    def __len__(self):
+        return len(self.index_mapper)
+
+    def get_raw_image(self, index, image_key="image"):
+        index, caption_index = self.index_mapper[index]
+        image_bytes = io.BytesIO(self.table[image_key][index].as_py())
+        image_bytes.seek(0)
+        return Image.open(image_bytes).convert("RGB")
+
+    def get_image(self, index, image_key="image"):
+        image = self.get_raw_image(index, image_key=image_key)
+        image_tensor = [tr(image) for tr in self.transforms]
+        return {
+            "image": image_tensor,
+            "img_index": self.index_mapper[index][0],
+            "cap_index": self.index_mapper[index][1],
+            "raw_index": index,
+        }
+
+    def get_false_image(self, rep, image_key="image"):
+        random_index = random.randint(0, len(self.index_mapper) - 1)
+        image = self.get_raw_image(random_index, image_key=image_key)
+        image_tensor = [tr(image) for tr in self.transforms]
+        return {f"false_image_{rep}": image_tensor}
+
+    def get_text(self, raw_index):
+        index, caption_index = self.index_mapper[raw_index]
+
+        text = self.all_texts[index][caption_index]
+        encoding = self.tokenizer(
+            text,
+            padding="max_length",
+            truncation=True,
+            max_length=self.max_text_len,
+            return_special_tokens_mask=True,
+        )
+        return {
+            "text": (text, encoding),
+            "img_index": index,
+            "cap_index": caption_index,
+            "raw_index": raw_index,
+        }
+
+    def get_false_text(self, rep):
+        random_index = random.randint(0, len(self.index_mapper) - 1)
+
+        index, caption_index = self.index_mapper[random_index]
+        text = self.all_texts[index][caption_index]
+        encoding = self.tokenizer(
+            text,
+            truncation=True,
+            max_length=self.max_text_len,
+            return_special_tokens_mask=True,
+        )
+        return {f"false_text_{rep}": (text, encoding)}
+
+    def get_suite(self, index):
+        result = None
+        while result is None:
+            try:
+                ret = dict()
+                ret.update(self.get_image(index))
+                if not self.image_only:
+                    txt = self.get_text(index)
+                    ret.update({"replica": True if txt["cap_index"] > 0 else False})
+                    ret.update(txt)
+
+                for i in range(self.draw_false_image):
+                    ret.update(self.get_false_image(i))
+                for i in range(self.draw_false_text):
+                    ret.update(self.get_false_text(i))
+                result = True
+            except Exception as e:
+                print(f"Error while read file idx {index} in {self.names[0]} -> {e}")
+                index = random.randint(0, len(self.index_mapper) - 1)
+        return ret
+
+
+    def get_text_suite(self, index):
+        result = None
+        while result is None:
+            try:
+                ret = dict()
+                txt = self.get_text(index)
+                ret.update({"replica": True if txt["cap_index"] > 0 else False})
+                ret.update(txt)
+                result = True
+            except Exception as e:
+                print(f"Error while read file idx {index} in {self.names[0]} -> {e}")
+                index = random.randint(0, len(self.index_mapper) - 1)
+        return ret
+
+
+    def collate(self, batch, mlm_collator):
+        batch_size = len(batch)
+        keys = set([key for b in batch for key in b.keys()])
+        dict_batch = {k: [dic[k] if k in dic else None for dic in batch] for k in keys}
+
+        img_keys = [k for k in list(dict_batch.keys()) if "image" in k]
+
+        for img_key in img_keys:
+            new_imgs = [tmp_img[0] for tmp_img in dict_batch[img_key]]
+            batch_new_imgs = torch.stack(new_imgs, dim=0)
+            dict_batch[img_key] = [batch_new_imgs]
+
+        txt_keys = [k for k in list(dict_batch.keys()) if "text" in k]
+
+        if len(txt_keys) != 0:
+            texts = [[d[0] for d in dict_batch[txt_key]] for txt_key in txt_keys]
+            encodings = [[d[1] for d in dict_batch[txt_key]] for txt_key in txt_keys]
+            draw_text_len = len(encodings)
+            flatten_encodings = [e for encoding in encodings for e in encoding]
+            flatten_mlms = mlm_collator(flatten_encodings)
+
+            for i, txt_key in enumerate(txt_keys):
+                texts, encodings = (
+                    [d[0] for d in dict_batch[txt_key]],
+                    [d[1] for d in dict_batch[txt_key]],
+                )
+
+                mlm_ids, mlm_labels = (
+                    flatten_mlms["input_ids"][batch_size * (i) : batch_size * (i + 1)],
+                    flatten_mlms["labels"][batch_size * (i) : batch_size * (i + 1)],
+                )
+
+                input_ids = torch.zeros_like(mlm_ids)
+                attention_mask = torch.zeros_like(mlm_ids)
+                for _i, encoding in enumerate(encodings):
+                    _input_ids, _attention_mask = (
+                        torch.tensor(encoding["input_ids"]),
+                        torch.tensor(encoding["attention_mask"]),
+                    )
+                    input_ids[_i, : len(_input_ids)] = _input_ids
+                    attention_mask[_i, : len(_attention_mask)] = _attention_mask
+
+                dict_batch[txt_key] = texts
+                dict_batch[f"{txt_key}_ids"] = input_ids
+                dict_batch[f"{txt_key}_labels"] = torch.full_like(input_ids, -100)
+                dict_batch[f"{txt_key}_ids_mlm"] = mlm_ids
+                dict_batch[f"{txt_key}_labels_mlm"] = mlm_labels
+                dict_batch[f"{txt_key}_masks"] = attention_mask
+
+        return dict_batch

vmlo/vlmo/datasets/coco_caption_karpathy_dataset.py

+from .base_dataset import BaseDataset
+
+
+class CocoCaptionKarpathyDataset(BaseDataset):
+    def __init__(self, *args, split="", **kwargs):
+        assert split in ["train", "val", "test"]
+        self.split = split
+
+        if split == "train":
+            names = ["coco_caption_karpathy_train", "coco_caption_karpathy_restval"]
+        elif split == "val":
+            names = ["coco_caption_karpathy_val"]
+        elif split == "test":
+            names = ["coco_caption_karpathy_test"]
+
+        super().__init__(*args, **kwargs, names=names, text_column_name="caption")
+
+    def __getitem__(self, index):
+        suite = self.get_suite(index)
+
+        if "test" in self.split:
+            _index, _question_index = self.index_mapper[index]
+            iid = self.table["image_id"][_index].as_py()
+            iid = int(iid.split(".")[0].split("_")[-1])
+            suite.update({"iid": iid})
+
+        return suite

vmlo/vlmo/datasets/conceptual_caption_dataset.py

+from glob import glob
+from .base_dataset import BaseDataset
+
+
+class ConceptualCaptionDataset(BaseDataset):
+    def __init__(self, *args, split="", **kwargs):
+        assert split in ["train", "val", "test"]
+        if split == "test":
+            split = "val"
+
+        if split == "train":
+            names = [f"conceptual_caption_train_{i}" for i in range(30)]
+        elif split == "val":
+            names = ["conceptual_caption_val_0"]
+
+        super().__init__(*args, **kwargs, names=names, text_column_name="caption")
+
+    def __getitem__(self, index):
+        return self.get_suite(index)

vmlo/vlmo/datasets/f30k_caption_karpathy_dataset.py

+from .base_dataset import BaseDataset
+
+
+class F30KCaptionKarpathyDataset(BaseDataset):
+    def __init__(self, *args, split="", **kwargs):
+        assert split in ["train", "val", "test"]
+
+        if split == "train":
+            names = ["f30k_caption_karpathy_train"]
+        elif split == "val":
+            names = ["f30k_caption_karpathy_val"]
+        elif split == "test":
+            names = ["f30k_caption_karpathy_test"]
+
+        super().__init__(*args, **kwargs, names=names, text_column_name="caption")
+
+    def __getitem__(self, index):
+        return self.get_suite(index)

vmlo/vlmo/datasets/nlvr2_dataset.py

+from .base_dataset import BaseDataset
+import sys
+import random
+
+
+class NLVR2Dataset(BaseDataset):
+    def __init__(self, *args, split="", **kwargs):
+        assert split in ["train", "val", "test"]
+        self.split = split
+
+        if split == "train":
+            names = ["nlvr2_train"]
+        elif split == "val":
+            names = ["nlvr2_dev", "nlvr2_test1"]
+        elif split == "test":
+            names = ["nlvr2_dev", "nlvr2_test1"]
+
+        super().__init__(
+            *args,
+            **kwargs,
+            names=names,
+            text_column_name="questions",
+            remove_duplicate=False,
+        )
+
+    def __getitem__(self, index):
+        result = None
+        while result is None:
+            try:
+                image_tensor_0 = self.get_image(index, image_key="image_0")["image"]
+                image_tensor_1 = self.get_image(index, image_key="image_1")["image"]
+                text = self.get_text(index)["text"]
+                result = True
+            except:
+                print(
+                    f"error while read file idx {index} in {self.names[0]}",
+                    file=sys.stderr,
+                )
+                index = random.randint(0, len(self.index_mapper) - 1)
+
+        index, question_index = self.index_mapper[index]
+        answers = self.table["answers"][index][question_index].as_py()
+        answers = answers == "True"
+
+        return {
+            "image_0": image_tensor_0,
+            "image_1": image_tensor_1,
+            "text": text,
+            "answers": answers,
+            "table_name": self.table_names[index],
+        }

vmlo/vlmo/datasets/sbu_caption_dataset.py

+from glob import glob
+from .base_dataset import BaseDataset
+
+
+class SBUCaptionDataset(BaseDataset):
+    def __init__(self, *args, split="", **kwargs):
+        assert split in ["train", "val", "test"]
+        if split == "test":
+            split = "val"
+
+        if split == "train":
+            names = [f"sbu_{i}" for i in range(9)]
+        elif split == "val":
+            names = []
+
+        super().__init__(*args, **kwargs, names=names, text_column_name="caption")
+
+    def __getitem__(self, index):
+        return self.get_suite(index)

vmlo/vlmo/datasets/vg_caption_dataset.py

+from .base_dataset import BaseDataset
+
+
+class VisualGenomeCaptionDataset(BaseDataset):
+    def __init__(self, *args, split="", **kwargs):
+        assert split in ["train", "val", "test"]
+        if split == "test":
+            split = "val"
+
+        if split == "train":
+            names = ["vg"]
+        elif split == "val":
+            names = []
+
+        super().__init__(*args, **kwargs, names=names, text_column_name="caption")
+
+    def __getitem__(self, index):
+        return self.get_suite(index)

vmlo/vlmo/datasets/vqav2_dataset.py

+from .base_dataset import BaseDataset
+
+
+class VQAv2Dataset(BaseDataset):
+    def __init__(self, *args, split="", **kwargs):
+        assert split in ["train", "val", "test"]
+        self.split = split
+
+        if split == "train":
+            names = ["vqav2_train", "vqav2_trainable_val"]
+        elif split == "val":
+            names = ["vqav2_rest_val"]
+        elif split == "test":
+            names = ["vqav2_test"]  # vqav2_test-dev for test-dev
+
+        super().__init__(
+            *args,
+            **kwargs,
+            names=names,
+            text_column_name="questions",
+            remove_duplicate=False,
+        )
+
+    def __getitem__(self, index):
+        image_tensor = self.get_image(index)["image"]
+        text = self.get_text(index)["text"]
+
+        index, question_index = self.index_mapper[index]
+        qid = self.table["question_id"][index][question_index].as_py()
+
+        if self.split != "test":
+            answers = self.table["answers"][index][question_index].as_py()
+            labels = self.table["answer_labels"][index][question_index].as_py()
+            scores = self.table["answer_scores"][index][question_index].as_py()
+        else:
+            answers = list()
+            labels = list()
+            scores = list()
+
+        return {
+            "image": image_tensor,
+            "text": text,
+            "vqa_answer": answers,
+            "vqa_labels": labels,
+            "vqa_scores": scores,
+            "qid": qid,
+        }

vmlo/vlmo/datasets/wikibk_dataset.py

+from glob import glob
+from .base_dataset import BaseDataset
+
+
+class WikibkDataset(BaseDataset):
+    def __init__(self, *args, split="", **kwargs):
+        assert split in ["train", "val", "test"]
+        if split == "test":
+            split = "val"
+
+        if split == "train":
+            names = [f"wikibk_train_{i}" for i in range(50)]
+        elif split == "val":
+            names = ["wikibk_val_0"]
+
+        super().__init__(*args, **kwargs, names=names, text_column_name="caption")
+
+    def __getitem__(self, index):
+        return self.get_text_suite(index)

vmlo/vlmo/gadgets/my_metrics.py

+import torch
+from torchmetrics import Metric
+
+
+class Accuracy(Metric):
+    def __init__(self, dist_sync_on_step=False):
+        super().__init__(dist_sync_on_step=dist_sync_on_step)
+        self.add_state("correct", default=torch.tensor(0.0), dist_reduce_fx="sum")
+        self.add_state("total", default=torch.tensor(0.0), dist_reduce_fx="sum")
+
+    def update(self, logits, target):
+        logits, target = (
+            logits.detach().to(self.correct.device),
+            target.detach().to(self.correct.device),
+        )
+        preds = logits.argmax(dim=-1)
+        preds = preds[target != -100]
+        target = target[target != -100]
+        if target.numel() == 0:
+            return 1
+
+        assert preds.shape == target.shape
+
+        self.correct += torch.sum(preds == target)
+        self.total += target.numel()
+
+    def compute(self):
+        return self.correct / self.total
+
+
+class Scalar(Metric):
+    def __init__(self, dist_sync_on_step=False):
+        super().__init__(dist_sync_on_step=dist_sync_on_step)
+        self.add_state("scalar", default=torch.tensor(0.0), dist_reduce_fx="sum")
+        self.add_state("total", default=torch.tensor(0.0), dist_reduce_fx="sum")
+
+    def update(self, scalar):
+        if isinstance(scalar, torch.Tensor):
+            scalar = scalar.detach().to(self.scalar.device)
+        else:
+            scalar = torch.tensor(scalar).float().to(self.scalar.device)
+        self.scalar += scalar
+        self.total += 1
+
+    def compute(self):
+        return self.scalar / self.total
+
+
+class VQAScore(Metric):
+    def __init__(self, dist_sync_on_step=False):
+        super().__init__(dist_sync_on_step=dist_sync_on_step)
+        self.add_state("score", default=torch.tensor(0.0), dist_reduce_fx="sum")
+        self.add_state("total", default=torch.tensor(0.0), dist_reduce_fx="sum")
+
+    def update(self, logits, target):
+        logits, target = (
+            logits.detach().float().to(self.score.device),
+            target.detach().float().to(self.score.device),
+        )
+        logits = torch.max(logits, 1)[1]
+        one_hots = torch.zeros(*target.size()).to(target)
+        one_hots.scatter_(1, logits.view(-1, 1), 1)
+        scores = one_hots * target
+
+        self.score += scores.sum()
+        self.total += len(logits)
+
+    def compute(self):
+        return self.score / self.total

vmlo/vlmo/modules/__init__.py

+from .vlmo_module import VLMo

vmlo/vlmo/modules/dist_utils.py

+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
+"""
+This file contains primitives for multi-gpu communication.
+This is useful when doing distributed training.
+"""
+
+import functools
+import logging
+import numpy as np
+import pickle
+import torch
+import torch.distributed as dist
+
+import torch
+
+_LOCAL_PROCESS_GROUP = None
+"""
+A torch process group which only includes processes that on the same machine as the current process.
+This variable is set when processes are spawned by `launch()` in "engine/launch.py".
+"""
+
+
+def get_world_size() -> int:
+    if not dist.is_available():
+        return 1
+    if not dist.is_initialized():
+        return 1
+    return dist.get_world_size()
+
+
+def get_rank() -> int:
+    if not dist.is_available():
+        return 0
+    if not dist.is_initialized():
+        return 0
+    return dist.get_rank()
+
+
+def get_local_rank() -> int:
+    """
+    Returns:
+        The rank of the current process within the local (per-machine) process group.
+    """
+    if not dist.is_available():
+        return 0
+    if not dist.is_initialized():
+        return 0
+    assert _LOCAL_PROCESS_GROUP is not None
+    return dist.get_rank(group=_LOCAL_PROCESS_GROUP)
+
+
+def get_local_size() -> int:
+    """
+    Returns:
+        The size of the per-machine process group,
+        i.e. the number of processes per machine.
+    """
+    if not dist.is_available():
+        return 1
+    if not dist.is_initialized():
+        return 1
+    return dist.get_world_size(group=_LOCAL_PROCESS_GROUP)
+
+
+def is_main_process() -> bool:
+    return get_rank() == 0
+
+
+def synchronize():
+    """
+    Helper function to synchronize (barrier) among all processes when
+    using distributed training
+    """
+    if not dist.is_available():
+        return
+    if not dist.is_initialized():
+        return
+    world_size = dist.get_world_size()
+    if world_size == 1:
+        return
+    dist.barrier()
+
+
+@functools.lru_cache()
+def _get_global_gloo_group():
+    """
+    Return a process group based on gloo backend, containing all the ranks
+    The result is cached.
+    """
+    if dist.get_backend() == "nccl":
+        return dist.new_group(backend="gloo")
+    else:
+        return dist.group.WORLD
+
+
+def _serialize_to_tensor(data, group):
+    backend = dist.get_backend(group)
+    assert backend in ["gloo", "nccl"]
+    device = torch.device("cpu" if backend == "gloo" else "cuda")
+
+    buffer = pickle.dumps(data)
+    if len(buffer) > 1024 ** 3:
+        logger = logging.getLogger(__name__)
+        logger.warning(
+            "Rank {} trying to all-gather {:.2f} GB of data on device {}".format(
+                get_rank(), len(buffer) / (1024 ** 3), device
+            )
+        )
+    storage = torch.ByteStorage.from_buffer(buffer)
+    tensor = torch.ByteTensor(storage).to(device=device)
+    return tensor
+
+
+def _pad_to_largest_tensor(tensor, group):
+    """
+    Returns:
+        list[int]: size of the tensor, on each rank
+        Tensor: padded tensor that has the max size
+    """
+    world_size = dist.get_world_size(group=group)
+    assert (
+        world_size >= 1
+    ), "comm.gather/all_gather must be called from ranks within the given group!"
+    local_size = torch.tensor([tensor.numel()], dtype=torch.int64, device=tensor.device)
+    size_list = [
+        torch.zeros([1], dtype=torch.int64, device=tensor.device)
+        for _ in range(world_size)
+    ]
+    dist.all_gather(size_list, local_size, group=group)
+    size_list = [int(size.item()) for size in size_list]
+
+    max_size = max(size_list)
+
+    # we pad the tensor because torch all_gather does not support
+    # gathering tensors of different shapes
+    if local_size != max_size:
+        padding = torch.zeros(
+            (max_size - local_size,), dtype=torch.uint8, device=tensor.device
+        )
+        tensor = torch.cat((tensor, padding), dim=0)
+    return size_list, tensor
+
+
+def all_gather(data, group=None):
+    """
+    Run all_gather on arbitrary picklable data (not necessarily tensors).
+
+    Args:
+        data: any picklable object
+        group: a torch process group. By default, will use a group which
+            contains all ranks on gloo backend.
+
+    Returns:
+        list[data]: list of data gathered from each rank
+    """
+    if get_world_size() == 1:
+        return [data]
+    if group is None:
+        group = _get_global_gloo_group()
+    if dist.get_world_size(group) == 1:
+        return [data]
+
+    tensor = _serialize_to_tensor(data, group)
+
+    size_list, tensor = _pad_to_largest_tensor(tensor, group)
+    max_size = max(size_list)
+
+    # receiving Tensor from all ranks
+    tensor_list = [
+        torch.empty((max_size,), dtype=torch.uint8, device=tensor.device)
+        for _ in size_list
+    ]
+    dist.all_gather(tensor_list, tensor, group=group)
+
+    data_list = []
+    for size, tensor in zip(size_list, tensor_list):
+        buffer = tensor.cpu().numpy().tobytes()[:size]
+        data_list.append(pickle.loads(buffer))
+
+    return data_list
+
+
+def gather(data, dst=0, group=None):
+    """
+    Run gather on arbitrary picklable data (not necessarily tensors).
+
+    Args:
+        data: any picklable object
+        dst (int): destination rank
+        group: a torch process group. By default, will use a group which
+            contains all ranks on gloo backend.
+
+    Returns:
+        list[data]: on dst, a list of data gathered from each rank. Otherwise,
+            an empty list.
+    """
+    if get_world_size() == 1:
+        return [data]
+    if group is None:
+        group = _get_global_gloo_group()
+    if dist.get_world_size(group=group) == 1:
+        return [data]
+    rank = dist.get_rank(group=group)
+
+    tensor = _serialize_to_tensor(data, group)
+    size_list, tensor = _pad_to_largest_tensor(tensor, group)
+
+    # receiving Tensor from all ranks
+    if rank == dst:
+        max_size = max(size_list)
+        tensor_list = [
+            torch.empty((max_size,), dtype=torch.uint8, device=tensor.device)
+            for _ in size_list
+        ]
+        dist.gather(tensor, tensor_list, dst=dst, group=group)
+
+        data_list = []
+        for size, tensor in zip(size_list, tensor_list):
+            buffer = tensor.cpu().numpy().tobytes()[:size]
+            data_list.append(pickle.loads(buffer))
+        return data_list
+    else:
+        dist.gather(tensor, [], dst=dst, group=group)
+        return []
+
+
+def shared_random_seed():
+    """
+    Returns:
+        int: a random number that is the same across all workers.
+            If workers need a shared RNG, they can use this shared seed to
+            create one.
+
+    All workers must call this function, otherwise it will deadlock.
+    """
+    ints = np.random.randint(2 ** 31)
+    all_ints = all_gather(ints)
+    return all_ints[0]
+
+
+def reduce_dict(input_dict, average=True):
+    """
+    Reduce the values in the dictionary from all processes so that process with rank
+    0 has the reduced results.
+
+    Args:
+        input_dict (dict): inputs to be reduced. All the values must be scalar CUDA Tensor.
+        average (bool): whether to do average or sum
+
+    Returns:
+        a dict with the same keys as input_dict, after reduction.
+    """
+    world_size = get_world_size()
+    if world_size < 2:
+        return input_dict
+    with torch.no_grad():
+        names = []
+        values = []
+        # sort the keys so that they are consistent across processes
+        for k in sorted(input_dict.keys()):
+            names.append(k)
+            values.append(input_dict[k])
+        values = torch.stack(values, dim=0)
+        dist.reduce(values, dst=0)
+        if dist.get_rank() == 0 and average:
+            # only main process gets accumulated, so only divide by
+            # world_size in this case
+            values /= world_size
+        reduced_dict = {k: v for k, v in zip(names, values)}
+    return reduced_dict

vmlo/vlmo/modules/heads.py

+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from transformers.models.bert.modeling_bert import BertPredictionHeadTransform
+
+
+class Pooler(nn.Module):
+    def __init__(self, hidden_size):
+        super().__init__()
+        self.dense = nn.Linear(hidden_size, hidden_size)
+        self.activation = nn.Tanh()
+
+    def forward(self, hidden_states):
+        first_token_tensor = hidden_states[:, 0]
+        pooled_output = self.dense(first_token_tensor)
+        pooled_output = self.activation(pooled_output)
+        return pooled_output
+
+
+class ITMHead(nn.Module):
+    def __init__(self, hidden_size):
+        super().__init__()
+        self.fc = nn.Linear(hidden_size, 2)
+
+    def forward(self, x):
+        x = self.fc(x)
+        return x
+
+class ITCHead(nn.Module):
+    def __init__(self, hidden_size):
+        super().__init__()
+        self.fc = nn.Linear(hidden_size, hidden_size, bias=False)
+
+    def forward(self, x):
+        x = self.fc(x)
+        return x
+
+
+class MLMHead(nn.Module):
+    def __init__(self, config, weight=None):
+        super().__init__()
+        self.transform = BertPredictionHeadTransform(config)
+        self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.bias = nn.Parameter(torch.zeros(config.vocab_size))
+        if weight is not None:
+            self.decoder.weight = weight
+
+    def forward(self, x):
+        x = self.transform(x)
+        x = self.decoder(x) + self.bias
+        return x

vmlo/vlmo/modules/multiway_transformer.py

+""" Vision Transformer (ViT) in PyTorch
+
+A PyTorch implement of Vision Transformers as described in
+'An Image Is Worth 16 x 16 Words: Transformers for Image Recognition at Scale' - https://arxiv.org/abs/2010.11929
+
+The official jax code is released and available at https://github.com/google-research/vision_transformer
+
+Acknowledgments:
+* The paper authors for releasing code and weights, thanks!
+* I fixed my class token impl based on Phil Wang's https://github.com/lucidrains/vit-pytorch ... check it out
+for some einops/einsum fun
+* Simple transformer style inspired by Andrej Karpathy's https://github.com/karpathy/minGPT
+* Bert reference code checks against Huggingface Transformers and Tensorflow Bert
+
+DeiT model defs and weights from https://github.com/facebookresearch/deit,
+paper `DeiT: Data-efficient Image Transformers` - https://arxiv.org/abs/2012.12877
+
+Hacked together by / Copyright 2020 Ross Wightman
+"""
+from functools import partial
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from functools import partial
+
+from timm.models.layers import DropPath, to_2tuple, trunc_normal_
+from timm.models.registry import register_model
+from pytorch_lightning.utilities.distributed import rank_zero_info
+
+
+class Mlp(nn.Module):
+    def __init__(
+        self,
+        in_features,
+        hidden_features=None,
+        out_features=None,
+        act_layer=nn.GELU,
+        drop=0.0,
+    ):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim,
+        num_heads=8,
+        qkv_bias=False,
+        qk_scale=None,
+        attn_drop=0.0,
+        proj_drop=0.0,
+    ):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        # NOTE scale factor was wrong in my original version, can set manually to be compat with prev weights
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=False)
+        if qkv_bias:
+            self.q_bias = nn.Parameter(torch.zeros(dim))
+            self.v_bias = nn.Parameter(torch.zeros(dim))
+        else:
+            self.q_bias = None
+            self.v_bias = None
+        
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+
+    def forward(self, x, mask=None, relative_position_bias=None):
+        B, N, C = x.shape
+
+        qkv_bias = None
+        if self.q_bias is not None:
+            qkv_bias = torch.cat((self.q_bias, torch.zeros_like(self.v_bias, requires_grad=False), self.v_bias))
+        qkv = F.linear(input=x, weight=self.qkv.weight, bias=qkv_bias)
+        qkv = qkv.reshape(B, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4)
+
+        q, k, v = (
+            qkv[0],
+            qkv[1],
+            qkv[2],
+        )  # make torchscript happy (cannot use tensor as tuple)
+
+        q = q * self.scale
+        attn = (q.float() @ k.float().transpose(-2, -1))
+        
+        if relative_position_bias is not None:
+            attn = attn + relative_position_bias.unsqueeze(0)
+
+        if mask is not None:
+            mask = mask.bool()
+            attn = attn.masked_fill(~mask[:, None, None, :], float("-inf"))
+        attn = attn.softmax(dim=-1).type_as(x)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x
+
+
+class Block(nn.Module):
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        mlp_ratio=4.0,
+        qkv_bias=False,
+        qk_scale=None,
+        drop=0.0,
+        attn_drop=0.0,
+        drop_path=0.0,
+        act_layer=nn.GELU,
+        norm_layer=nn.LayerNorm,
+        with_vlffn=False,
+        layer_scale_init_values=0.1,
+        max_text_len=40,
+    ):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            qk_scale=qk_scale,
+            attn_drop=attn_drop,
+            proj_drop=drop,
+        )
+        # NOTE: drop path for stochastic depth, we shall see if this is better than dropout here
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+        self.norm2_text = norm_layer(dim)
+        self.norm2_imag = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp_text = Mlp(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop,
+        )
+        self.mlp_imag = Mlp(
+            in_features=dim,
+            hidden_features=mlp_hidden_dim,
+            act_layer=act_layer,
+            drop=drop,
+        )
+        self.mlp_vl = None
+        if with_vlffn:
+            self.mlp_vl = Mlp(
+                in_features=dim,
+                hidden_features=mlp_hidden_dim,
+                act_layer=act_layer,
+                drop=drop,
+            )
+            self.norm2_vl = norm_layer(dim)
+        
+        self.gamma_1 = \
+            nn.Parameter(layer_scale_init_values * torch.ones((dim)),requires_grad=True) \
+            if layer_scale_init_values is not None else 1.0
+        self.gamma_2 = \
+            nn.Parameter(layer_scale_init_values * torch.ones((dim)),requires_grad=True) \
+            if layer_scale_init_values is not None else 1.0
+
+        self.max_text_len = max_text_len
+
+    def forward(self, x, mask=None, modality_type=None, relative_position_bias=None):
+        x = x + self.drop_path(self.gamma_1 * self.attn(self.norm1(x), mask=mask, relative_position_bias=relative_position_bias))
+
+        if modality_type == "image":
+            x = x + self.drop_path(self.gamma_2 * self.mlp_imag(self.norm2_imag(x)))
+        elif modality_type == "text":
+            x = x + self.drop_path(self.gamma_2 * self.mlp_text(self.norm2_text(x)))
+        else:
+            if self.mlp_vl is None:
+                x_text = x[:, : self.max_text_len]
+                x_imag = x[:, self.max_text_len :]
+                x_text = x_text + self.drop_path(self.gamma_2 * self.mlp_text(self.norm2_text(x_text)))
+                x_imag = x_imag + self.drop_path(self.gamma_2 * self.mlp_imag(self.norm2_imag(x_imag)))
+                x = torch.cat([x_text, x_imag], dim=1)
+            else:
+                x = x + self.drop_path(self.gamma_2 * self.mlp_vl(self.norm2_vl(x)))
+
+        return x
+
+
+class PatchEmbed(nn.Module):
+    """ Image to Patch Embedding"""
+
+    def __init__(
+        self,
+        img_size=224,
+        patch_size=16,
+        in_chans=3,
+        embed_dim=768,
+        no_patch_embed_bias=False,
+    ):
+        super().__init__()
+        img_size = to_2tuple(img_size)
+        patch_size = to_2tuple(patch_size)
+        num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0])
+        self.patch_shape = (img_size[0] // patch_size[0], img_size[1] // patch_size[1])
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = num_patches
+
+        self.proj = nn.Conv2d(
+            in_chans,
+            embed_dim,
+            kernel_size=patch_size,
+            stride=patch_size,
+            bias=False if no_patch_embed_bias else True,
+        )
+
+    def forward(self, x):
+        B, C, H, W = x.shape
+        assert H == self.img_size[0] and W == self.img_size[1], \
+            f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})."
+        # FIXME look at relaxing size constraints
+        x = self.proj(x)
+        return x
+
+
+class MultiWayTransformer(nn.Module):
+    """ Vision Transformer
+
+    A PyTorch impl of : `An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale`  -
+        https://arxiv.org/abs/2010.11929
+    """
+
+    def __init__(
+        self,
+        img_size=224,
+        patch_size=16,
+        in_chans=3,
+        embed_dim=768,
+        depth=12,
+        num_heads=12,
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        qk_scale=None,
+        drop_rate=0.0,
+        attn_drop_rate=0.0,
+        drop_path_rate=0.0,
+        norm_layer=None,
+        need_relative_position_embed=True,
+        use_abs_pos_emb=False,
+        layer_scale_init_values=0.1,
+        vlffn_start_layer_index=10,
+        config=None,
+    ):
+        """
+        Args:
+            img_size (int, tuple): input image size
+            patch_size (int, tuple): patch size
+            in_chans (int): number of input channels
+            num_classes (int): number of classes for classification head
+            embed_dim (int): embedding dimension
+            depth (int): depth of transformer
+            num_heads (int): number of attention heads
+            mlp_ratio (int): ratio of mlp hidden dim to embedding dim
+            qkv_bias (bool): enable bias for qkv if True
+            qk_scale (float): override default qk scale of head_dim ** -0.5 if set
+            drop_rate (float): dropout rate
+            attn_drop_rate (float): attention dropout rate
+            drop_path_rate (float): stochastic depth rate
+            norm_layer: (nn.Module): normalization layer
+            need_relative_position_embed (bool): enable relative position bias on self-attention
+            use_abs_pos_emb (bool): enable abs pos emb
+            layer_scale_init_values (float or None): layer scale init values, set None to disable
+            vlffn_start_layer_index (int): vl-ffn start index
+            config: (dict): other hyper from pytorch-lighting
+        """
+        super().__init__()
+        drop_path_rate = drop_path_rate if config is None else config["drop_path_rate"]
+        rank_zero_info("drop path rate: {}".format(drop_path_rate))
+        self.use_abs_pos_emb = use_abs_pos_emb
+        self.need_relative_position_embed = need_relative_position_embed
+
+        self.num_features = (
+            self.embed_dim
+        ) = embed_dim  # num_features for consistency with other models
+        norm_layer = norm_layer or partial(nn.LayerNorm, eps=1e-6)
+
+        self.patch_embed = PatchEmbed(
+            img_size=img_size,
+            patch_size=patch_size,
+            in_chans=in_chans,
+            embed_dim=embed_dim,
+        )
+        num_patches = self.patch_embed.num_patches
+        self.patch_size = patch_size
+        self.num_heads = num_heads
+        self.vlffn_start_layer_index = vlffn_start_layer_index
+        if config["loss_names"]["textmlm"] > 0:
+            self.vlffn_start_layer_index = depth
+            rank_zero_info("Set vlffn_start_layer_index={} for text-only pretraining".format(self.vlffn_start_layer_index))
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim)) if self.use_abs_pos_emb else None
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        dpr = [
+            x.item() for x in torch.linspace(0, drop_path_rate, depth)
+        ]  # stochastic depth decay rule
+        self.blocks = nn.ModuleList(
+            [
+                Block(
+                    dim=embed_dim,
+                    num_heads=num_heads,
+                    mlp_ratio=mlp_ratio,
+                    qkv_bias=qkv_bias,
+                    qk_scale=qk_scale,
+                    drop=drop_rate,
+                    attn_drop=attn_drop_rate,
+                    drop_path=dpr[i],
+                    norm_layer=norm_layer,
+                    with_vlffn=(i >= self.vlffn_start_layer_index),
+                    layer_scale_init_values=layer_scale_init_values,
+                    max_text_len=config["max_text_len"],
+                )
+                for i in range(depth)
+            ]
+        )
+        self.norm = norm_layer(embed_dim)
+
+        if self.pos_embed is not None:
+            trunc_normal_(self.pos_embed, std=0.02)
+        trunc_normal_(self.cls_token, std=0.02)
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    @torch.jit.ignore
+    def no_weight_decay(self):
+        return {"pos_embed", "cls_token"}
+
+    def visual_embed(self, _x):
+        x = self.patch_embed(_x)
+        x = x.flatten(2).transpose(1, 2)
+        B, L, _ = x.shape
+
+        cls_tokens = self.cls_token.expand(B, -1, -1)
+        x = torch.cat((cls_tokens, x), dim=1)
+        
+        if self.pos_embed is not None:
+            x = x + self.pos_embed
+        x = self.pos_drop(x)
+
+        x_mask = torch.ones(x.shape[0], x.shape[1])
+
+        return x, x_mask
+
+
+# VLMo base/p16
+@register_model
+def vlmo_base_patch16(pretrained=False, **kwargs):
+    img_size = kwargs.pop("img_size", 224)
+    model = MultiWayTransformer(
+        img_size=img_size, patch_size=16, embed_dim=768, depth=12, num_heads=12, 
+        mlp_ratio=4, qkv_bias=True, vlffn_start_layer_index=10, 
+        norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+# VLMo large/p16
+@register_model
+def vlmo_large_patch16(pretrained=False, **kwargs):
+    img_size = kwargs.pop("img_size", 224)
+    model = MultiWayTransformer(
+        img_size=img_size, patch_size=16, embed_dim=1024, depth=24, num_heads=16, 
+        mlp_ratio=4, qkv_bias=True, vlffn_start_layer_index=21, 
+        norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model
+
+# VLMo base+/p16
+@register_model
+def vlmo_base_plus_patch16(pretrained=False, **kwargs):
+    img_size = kwargs.pop("img_size", 224)
+    model = MultiWayTransformer(
+        img_size=img_size, patch_size=16, embed_dim=544, depth=24, num_heads=16, 
+        mlp_ratio=4, qkv_bias=True, vlffn_start_layer_index=21,
+        use_abs_pos_emb=True, need_relative_position_embed=False, 
+        layer_scale_init_values=None, norm_layer=partial(nn.LayerNorm, eps=1e-6), **kwargs)
+    return model

vmlo/vlmo/modules/objectives.py

+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import os
+import glob
+import json
+import tqdm
+import functools
+import torch.distributed as dist
+
+from torch.utils.data.distributed import DistributedSampler
+from einops import rearrange
+from pytorch_lightning.utilities.distributed import rank_zero_info
+
+from vlmo.modules.dist_utils import all_gather
+
+
+def compute_mlm(pl_module, batch):
+    infer = pl_module.infer(batch, mask_text=True, mask_image=False)
+    mlm_logits = pl_module.mlm_score(infer["text_feats"])
+    mlm_labels = infer["text_labels"]
+
+    mlm_loss = F.cross_entropy(
+        mlm_logits.view(-1, pl_module.hparams.config["vocab_size"]),
+        mlm_labels.view(-1),
+        ignore_index=-100,
+    )
+
+    ret = {
+        "mlm_loss": mlm_loss * 0.25,
+        "mlm_logits": mlm_logits,
+        "mlm_labels": mlm_labels,
+        "mlm_ids": infer["text_ids"],
+    }
+
+    phase = "train" if pl_module.training else "val"
+    loss = getattr(pl_module, f"{phase}_mlm_loss")(ret["mlm_loss"])
+    acc = getattr(pl_module, f"{phase}_mlm_accuracy")(
+        ret["mlm_logits"], ret["mlm_labels"]
+    )
+    pl_module.log(f"mlm/{phase}/loss", loss)
+    pl_module.log(f"mlm/{phase}/accuracy", acc)
+
+    return ret
+
+
+def compute_textonly_mlm(pl_module, batch):
+    infer = pl_module.infer_text_mlm(batch, mask_text=True)
+    mlm_logits = pl_module.mlm_score(infer["text_feats"])
+    mlm_labels = infer["text_labels"]
+
+    mlm_loss = F.cross_entropy(
+        mlm_logits.view(-1, pl_module.hparams.config["vocab_size"]),
+        mlm_labels.view(-1),
+        ignore_index=-100,
+    )
+
+    ret = {
+        "mlm_loss": mlm_loss,
+        "mlm_logits": mlm_logits,
+        "mlm_labels": mlm_labels,
+        "mlm_ids": infer["text_ids"],
+    }
+
+    phase = "train" if pl_module.training else "val"
+    loss = getattr(pl_module, f"{phase}_textmlm_loss")(ret["mlm_loss"])
+    acc = getattr(pl_module, f"{phase}_textmlm_accuracy")(
+        ret["mlm_logits"], ret["mlm_labels"]
+    )
+    pl_module.log(f"textmlm/{phase}/loss", loss)
+    pl_module.log(f"textmlm/{phase}/accuracy", acc)
+
+    return ret
+
+
+def compute_itm_hardneg(pl_module, batch, sim_i2t, sim_t2i):
+    pos_len = batch["text_ids"].size(0)
+    neg_len = batch["text_ids"].size(0)
+    bsz = batch["text_ids"].size(0)
+    itm_labels = torch.cat([torch.ones(pos_len), torch.zeros(neg_len), torch.zeros(neg_len)]).to(
+        pl_module.device
+    )
+
+    batch = {k: v for k, v in batch.items()}
+    infer_pos = pl_module.infer(batch, mask_text=False, mask_image=False)
+
+    batch_text_ids = infer_pos["text_ids"]
+    batch_text_masks = infer_pos["text_masks"]
+    batch_image = infer_pos["image"]
+
+    with torch.no_grad():
+        world_size = dist.get_world_size()
+        rank = dist.get_rank()
+
+        # We gather tensors from all gpus to get more hard negative candidates.
+        gathered_text_ids = [
+            torch.zeros_like(batch_text_ids) for _ in range(world_size)
+        ]
+        gathered_text_masks = [
+            torch.zeros_like(batch_text_masks) for _ in range(world_size)
+        ]
+        gathered_image = [
+            torch.zeros_like(batch_image) for _ in range(world_size)
+        ]
+
+        dist.all_gather(gathered_text_ids, batch_text_ids)
+        dist.all_gather(gathered_text_masks, batch_text_masks)
+        dist.all_gather(gathered_image, batch_image)
+
+        all_text_ids = torch.cat(
+            [batch_text_ids]
+            + gathered_text_ids[:rank]
+            + gathered_text_ids[rank + 1 :]
+        )
+        all_text_masks = torch.cat(
+            [batch_text_masks]
+            + gathered_text_masks[:rank]
+            + gathered_text_masks[rank + 1 :]
+        )
+        all_image = torch.cat(
+            [batch_image]
+            + gathered_image[:rank]
+            + gathered_image[rank + 1 :]
+        )
+
+    with torch.no_grad():       
+        weights_i2t = F.softmax(sim_i2t[:bsz, :].float(), dim=1)
+        weights_t2i = F.softmax(sim_t2i[:bsz, :].float(), dim=1)
+
+        weights_i2t.fill_diagonal_(0)
+        weights_t2i.fill_diagonal_(0)
+    
+    images_neg = []    
+    for b in range(bsz):
+        neg_idx = torch.multinomial(weights_t2i[b], 1).item()
+        images_neg.append(all_image[neg_idx])
+    images_neg = torch.stack(images_neg, dim=0)   
+
+    # select a negative text for each image
+    text_ids_neg = []
+    text_masks_neg = []
+    for b in range(bsz):
+        neg_idx = torch.multinomial(weights_i2t[b], 1).item()
+        text_ids_neg.append(all_text_ids[neg_idx])
+        text_masks_neg.append(all_text_masks[neg_idx])
+
+    text_ids_neg = torch.stack(text_ids_neg, dim=0)     
+    text_masks_neg = torch.stack(text_masks_neg, dim=0)      
+
+    # text_labels is not used in ITM loss
+    batch_imgs_neg = {"image":[images_neg], "text_ids":batch["text_ids"], "text_labels":batch["text_labels"], "text_masks":batch["text_masks"]}
+    infer_imags_neg = pl_module.infer(batch_imgs_neg, mask_text=False, mask_image=False)
+    
+    batch_text_neg = {"image":batch["image"], "text_ids":text_ids_neg, "text_labels":batch["text_labels"], "text_masks":text_masks_neg}
+    infer_text_neg = pl_module.infer(batch_text_neg, mask_text=False, mask_image=False)
+
+    all_cls_feats = torch.cat([infer_pos["cls_feats"], infer_imags_neg["cls_feats"], infer_text_neg["cls_feats"]], dim=0)
+
+    itm_logits = pl_module.itm_score(all_cls_feats)
+    itm_loss = F.cross_entropy(itm_logits, itm_labels.long())
+
+    ret = {
+        "itm_loss": itm_loss,
+        "itm_logits": itm_logits,
+        "itm_labels": itm_labels,
+    }
+
+    phase = "train" if pl_module.training else "val"
+    loss = getattr(pl_module, f"{phase}_itm_loss")(ret["itm_loss"])
+    acc = getattr(pl_module, f"{phase}_itm_accuracy")(
+        ret["itm_logits"], ret["itm_labels"]
+    )
+    pl_module.log(f"itm/{phase}/loss", loss)
+    pl_module.log(f"itm/{phase}/accuracy", acc)
+
+    return ret
+
+
+# The implementation of image-text contrastive refers to open_clip (https://github.com/mlfoundations/open_clip)
+def compute_itc(pl_module, batch, aggregate=True):
+    # pl_module.logit_scale.data = torch.clamp(pl_module.logit_scale.data, 0, 4.6052)
+
+    infer_imag = pl_module.infer_image(batch, mask_image=False)
+    infer_text = pl_module.infer_text(batch, mask_text=False)
+
+    image_features = infer_imag["cls_feats"]
+    text_features = infer_text["cls_feats"]
+    logit_scale = pl_module.logit_scale.exp().mean()
+
+    image_vlffn_features = infer_imag["cls_vlffn_feats"]
+    text_vlffn_features = infer_text["cls_vlffn_feats"]
+    logit_vl_scale = pl_module.logit_vl_scale.exp().mean()
+
+    if aggregate:
+        world_size = dist.get_world_size()
+        rank = dist.get_rank()
+
+        # We gather tensors from all gpus to get more negatives to contrast with.
+        gathered_image_features = [
+            torch.zeros_like(image_features) for _ in range(world_size)
+        ]
+        gathered_text_features = [
+            torch.zeros_like(text_features) for _ in range(world_size)
+        ]
+        dist.all_gather(gathered_image_features, image_features)
+        dist.all_gather(gathered_text_features, text_features)
+
+        all_image_features = torch.cat(
+            [image_features]
+            + gathered_image_features[:rank]
+            + gathered_image_features[rank + 1 :]
+        )
+        all_text_features = torch.cat(
+            [text_features]
+            + gathered_text_features[:rank]
+            + gathered_text_features[rank + 1 :]
+        )
+
+        # this is needed to send gradients back everywhere.
+        logits_per_image = logit_scale * all_image_features @ all_text_features.t()
+        logits_per_text = logits_per_image.t()
+
+        gathered_image_vlffn_features = [
+            torch.zeros_like(image_vlffn_features) for _ in range(world_size)
+        ]
+        gathered_text_vlffn_features = [
+            torch.zeros_like(text_vlffn_features) for _ in range(world_size)
+        ]
+        dist.all_gather(gathered_image_vlffn_features, image_vlffn_features)
+        dist.all_gather(gathered_text_vlffn_features, text_vlffn_features)
+
+        all_image_vlffn_features = torch.cat(
+            [image_vlffn_features]
+            + gathered_image_vlffn_features[:rank]
+            + gathered_image_vlffn_features[rank + 1 :]
+        )
+        all_text_vlffn_features = torch.cat(
+            [text_vlffn_features]
+            + gathered_text_vlffn_features[:rank]
+            + gathered_text_vlffn_features[rank + 1 :]
+        )
+        # this is needed to send gradients back everywhere.
+        logits_per_vlffn_image = logit_vl_scale * all_image_vlffn_features @ all_text_vlffn_features.t()
+        logits_per_vlffn_text = logits_per_vlffn_image.t()
+
+    else:
+        logits_per_image = logit_scale * image_features @ text_features.t()
+        logits_per_text = logit_scale * text_features @ image_features.t()
+
+    ground_truth = torch.arange(len(logits_per_image)).long().to(device=logits_per_image.get_device())
+
+    itc_loss = (
+        F.cross_entropy(logits_per_image.float(), ground_truth)
+        + F.cross_entropy(logits_per_text.float(), ground_truth)
+    ) / 2
+
+    itc_vlffn_loss = (
+        F.cross_entropy(logits_per_vlffn_image.float(), ground_truth)
+        + F.cross_entropy(logits_per_vlffn_text.float(), ground_truth)
+    ) / 2
+
+    itc_total_loss = (itc_loss + itc_vlffn_loss) * 0.5
+
+    ret = {
+        "itc_loss": itc_total_loss,
+        "itc_i2t_logits": logits_per_image,
+        "itc_t2i_logits": logits_per_text,
+        "itc_labels": ground_truth,
+        "itc_logit_scale": logit_scale,
+        "itc_logit_vl_scale": logit_vl_scale,
+    }
+
+    phase = "train" if pl_module.training else "val"
+    loss = getattr(pl_module, f"{phase}_itc_loss")(ret["itc_loss"])
+    scale = getattr(pl_module, f"{phase}_itc_logit_scale")(ret["itc_logit_scale"])
+    i2t_acc = getattr(pl_module, f"{phase}_itc_i2t_accuracy")(
+        ret["itc_i2t_logits"], ret["itc_labels"]
+    )
+    t2i_acc = getattr(pl_module, f"{phase}_itc_t2i_accuracy")(
+        ret["itc_t2i_logits"], ret["itc_labels"]
+    )
+    pl_module.log(f"itc/{phase}/loss", loss)
+    pl_module.log(f"itc/{phase}/logit_scale", scale)
+    pl_module.log(f"itc/{phase}/i2t_accuracy", i2t_acc)
+    pl_module.log(f"itc/{phase}/t2i_accuracy", t2i_acc)
+
+    vl_scale = getattr(pl_module, f"{phase}_itc_vl_logit_scale")(ret["itc_logit_vl_scale"])
+    vl_i2t_acc = getattr(pl_module, f"{phase}_itc_vl_i2t_accuracy")(
+        logits_per_vlffn_image, ret["itc_labels"]
+    )
+    vl_t2i_acc = getattr(pl_module, f"{phase}_itc_vl_t2i_accuracy")(
+        logits_per_vlffn_text, ret["itc_labels"]
+    )
+    pl_module.log(f"itc/{phase}/vl_logit_scale", vl_scale)
+    pl_module.log(f"itc/{phase}/vl_i2t_accuracy", vl_i2t_acc)
+    pl_module.log(f"itc/{phase}/vl_t2i_accuracy", vl_t2i_acc)
+
+    return ret
+
+
+def compute_irtr(pl_module, batch, aggregate=True):
+    # pl_module.logit_scale.data = torch.clamp(pl_module.logit_scale.data, 0, 4.6052)
+
+    infer_imag = pl_module.infer_image_ft(batch, mask_image=False)
+    infer_text = pl_module.infer_text_ft(batch, mask_text=False)
+
+    image_features = infer_imag["cls_feats"]
+    text_features = infer_text["cls_feats"]
+    logit_scale = pl_module.logit_scale.exp().mean()
+
+    if aggregate:
+        world_size = dist.get_world_size()
+        rank = dist.get_rank()
+
+        # We gather tensors from all gpus to get more negatives to contrast with.
+        gathered_image_features = [
+            torch.zeros_like(image_features) for _ in range(world_size)
+        ]
+        gathered_text_features = [
+            torch.zeros_like(text_features) for _ in range(world_size)
+        ]
+        dist.all_gather(gathered_image_features, image_features)
+        dist.all_gather(gathered_text_features, text_features)
+
+        all_image_features = torch.cat(
+            [image_features]
+            + gathered_image_features[:rank]
+            + gathered_image_features[rank + 1 :]
+        )
+        all_text_features = torch.cat(
+            [text_features]
+            + gathered_text_features[:rank]
+            + gathered_text_features[rank + 1 :]
+        )
+
+        # this is needed to send gradients back everywhere.
+        logits_per_image = logit_scale * all_image_features @ all_text_features.t()
+        logits_per_text = logits_per_image.t()
+    else:
+        logits_per_image = logit_scale * image_features @ text_features.t()
+        logits_per_text = logit_scale * text_features @ image_features.t()
+
+    ground_truth = torch.arange(len(logits_per_image)).long().to(device=logits_per_image.get_device())
+
+    irtr_loss = (
+        F.cross_entropy(logits_per_image.float(), ground_truth)
+        + F.cross_entropy(logits_per_text.float(), ground_truth)
+    ) / 2
+
+    ret = {
+        "irtr_loss": irtr_loss,
+        "irtr_i2t_logits": logits_per_image,
+        "irtr_t2i_logits": logits_per_text,
+        "irtr_labels": ground_truth,
+        "irtr_logit_scale": logit_scale,
+    }
+
+    phase = "train" if pl_module.training else "val"
+    loss = getattr(pl_module, f"{phase}_irtr_loss")(ret["irtr_loss"])
+    scale = getattr(pl_module, f"{phase}_irtr_logit_scale")(ret["irtr_logit_scale"])
+    i2t_acc = getattr(pl_module, f"{phase}_irtr_i2t_accuracy")(
+        ret["irtr_i2t_logits"], ret["irtr_labels"]
+    )
+    t2i_acc = getattr(pl_module, f"{phase}_irtr_t2i_accuracy")(
+        ret["irtr_t2i_logits"], ret["irtr_labels"]
+    )
+    pl_module.log(f"irtr/{phase}/loss", loss)
+    pl_module.log(f"irtr/{phase}/logit_scale", scale)
+    pl_module.log(f"irtr/{phase}/i2t_accuracy", i2t_acc)
+    pl_module.log(f"irtr/{phase}/t2i_accuracy", t2i_acc)
+
+    return ret
+
+
+def compute_vqa(pl_module, batch):
+    infer = pl_module.infer(batch, mask_text=False, mask_image=False)
+    vqa_logits = pl_module.vqa_classifier(infer["cls_feats"])
+    vqa_targets = torch.zeros(
+        len(vqa_logits), pl_module.hparams.config["vqav2_label_size"]
+    ).to(pl_module.device)
+
+    vqa_labels = batch["vqa_labels"]
+    vqa_scores = batch["vqa_scores"]
+
+    for i, (_label, _score) in enumerate(zip(vqa_labels, vqa_scores)):
+        for l, s in zip(_label, _score):
+            vqa_targets[i, l] = s
+
+    vqa_loss = (
+        F.binary_cross_entropy_with_logits(vqa_logits, vqa_targets)
+        * vqa_targets.shape[1]
+    )  # https://github.com/jnhwkim/ban-vqa/blob/master/train.py#L19
+
+    ret = {
+        "vqa_loss": vqa_loss,
+        "vqa_logits": vqa_logits,
+        "vqa_targets": vqa_targets,
+        "vqa_labels": vqa_labels,
+        "vqa_scores": vqa_scores,
+    }
+
+    phase = "train" if pl_module.training else "val"
+    loss = getattr(pl_module, f"{phase}_vqa_loss")(ret["vqa_loss"])
+    score = getattr(pl_module, f"{phase}_vqa_score")(
+        ret["vqa_logits"], ret["vqa_targets"]
+    )
+    pl_module.log(f"vqa/{phase}/loss", loss)
+    pl_module.log(f"vqa/{phase}/score", score)
+
+    return ret
+
+
+def compute_nlvr2(pl_module, batch):
+    infer1 = pl_module.infer(
+        batch, mask_text=False, mask_image=False, image_token_type_idx=1
+    )
+    infer2 = pl_module.infer(
+        batch, mask_text=False, mask_image=False, image_token_type_idx=2
+    )
+
+    cls_feats = torch.cat([infer1["cls_feats"], infer2["cls_feats"]], dim=-1)
+    nlvr2_logits = pl_module.nlvr2_classifier(cls_feats)
+
+    nlvr2_labels = batch["answers"]
+    nlvr2_labels = torch.tensor(nlvr2_labels).to(pl_module.device).long()
+    nlvr2_loss = F.cross_entropy(nlvr2_logits, nlvr2_labels)
+
+    ret = {
+        "nlvr2_loss": nlvr2_loss,
+        "nlvr2_logits": nlvr2_logits,
+        "nlvr2_labels": nlvr2_labels,
+    }
+
+    phase = "train" if pl_module.training else "val"
+
+    if phase == "train":
+        loss = getattr(pl_module, f"{phase}_nlvr2_loss")(ret["nlvr2_loss"])
+        acc = getattr(pl_module, f"{phase}_nlvr2_accuracy")(
+            ret["nlvr2_logits"], ret["nlvr2_labels"]
+        )
+        pl_module.log(f"nlvr2/{phase}/loss", loss)
+        pl_module.log(f"nlvr2/{phase}/accuracy", acc)
+    else:
+        dev_batches = [i for i, n in enumerate(batch["table_name"]) if "dev" in n]
+        test_batches = [i for i, n in enumerate(batch["table_name"]) if "test" in n]
+
+        if dev_batches:
+            dev_loss = getattr(pl_module, f"dev_nlvr2_loss")(
+                F.cross_entropy(
+                    ret["nlvr2_logits"][dev_batches], ret["nlvr2_labels"][dev_batches]
+                )
+            )
+            dev_acc = getattr(pl_module, f"dev_nlvr2_accuracy")(
+                ret["nlvr2_logits"][dev_batches], ret["nlvr2_labels"][dev_batches]
+            )
+            pl_module.log(f"nlvr2/dev/loss", dev_loss)
+            pl_module.log(f"nlvr2/dev/accuracy", dev_acc)
+        if test_batches:
+            test_loss = getattr(pl_module, f"test_nlvr2_loss")(
+                F.cross_entropy(
+                    ret["nlvr2_logits"][test_batches], ret["nlvr2_labels"][test_batches]
+                )
+            )
+            test_acc = getattr(pl_module, f"test_nlvr2_accuracy")(
+                ret["nlvr2_logits"][test_batches], ret["nlvr2_labels"][test_batches]
+            )
+            pl_module.log(f"nlvr2/test/loss", test_loss)
+            pl_module.log(f"nlvr2/test/accuracy", test_acc)
+
+    return ret
+
+
+@torch.no_grad()
+def compute_irtr_recall(pl_module, split="test"):
+    world_size = dist.get_world_size()
+    rank = dist.get_rank()
+    
+    if split == "val":
+        rank_zero_info("Use val set...")
+        text_dset = pl_module.trainer.datamodule.dms[0].make_no_false_val_dset()
+    else:
+        rank_zero_info("Use test set...")
+        text_dset = pl_module.trainer.datamodule.dms[0].make_no_false_test_dset()
+
+    text_dset.tokenizer = pl_module.trainer.datamodule.dms[0].tokenizer
+    text_loader = torch.utils.data.DataLoader(
+        text_dset,
+        batch_size=32,
+        num_workers=2, #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,
+        ),
+    )
+
+    if split == "val":
+        image_dset = pl_module.trainer.datamodule.dms[0].make_no_false_val_dset(
+            image_only=True
+        )
+    else:
+        image_dset = pl_module.trainer.datamodule.dms[0].make_no_false_test_dset(
+            image_only=True
+        )
+
+    image_dset.tokenizer = pl_module.trainer.datamodule.dms[0].tokenizer
+    image_loader = torch.utils.data.DataLoader(
+        image_dset,
+        batch_size=32,
+        num_workers=2, #pl_module.hparams.config["num_workers"],
+        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)
+
+    rank_zero_info("len(tiids): {}".format(len(tiids)))
+
+    image_preload = list()
+    for _b in tqdm.tqdm(image_loader, desc="image prefetch loop"):
+        image_preload.append(
+            {
+                "image": [_b["image"][0].to(pl_module.device)],
+                "img_index": _b["img_index"],
+            }
+        )
+    iids = list()
+    for pre in image_preload:
+        iids += pre["img_index"]
+    iids = torch.tensor(iids)
+
+    rank_zero_info("len(iids): {}".format(len(iids)))
+
+    txt_cls_feats = list()
+    for txt_batch in text_preload:
+        with torch.cuda.amp.autocast():
+            cls_feats = pl_module.infer_text_ft(
+                    {
+                        "text_ids": txt_batch["text_ids"],
+                        "text_masks": txt_batch["text_masks"],
+                        "text_labels": txt_batch["text_labels"],
+                    }
+                )["cls_feats"]
+        txt_cls_feats.append(cls_feats)
+
+    img_cls_feats = list()
+    for img_batch in image_preload:
+        with torch.cuda.amp.autocast():
+            cls_feats = pl_module.infer_image_ft(
+                    {
+                        "image": img_batch["image"],
+                    }
+                )["cls_feats"]
+        img_cls_feats.append(cls_feats)
+
+    txt_cls_feats = torch.cat(txt_cls_feats)
+    img_cls_feats = torch.cat(img_cls_feats)
+
+    rank_zero_info("txt_cls_feats.size(): {}\t{}".format(txt_cls_feats.size(), split))
+    rank_zero_info("img_cls_feats.size(): {}\t{}".format(img_cls_feats.size(), split))
+
+    scores = img_cls_feats @ txt_cls_feats.t()
+
+    rank_zero_info("scores.size(): {}".format(scores.size(), split))
+
+    topk10 = scores.topk(10, dim=1)
+    topk5 = scores.topk(5, dim=1)
+    topk1 = scores.topk(1, dim=1)
+    topk10_iids = tiids[topk10.indices.to(tiids.device)]
+    topk5_iids = tiids[topk5.indices.to(tiids.device)]
+    topk1_iids = tiids[topk1.indices.to(tiids.device)]
+
+    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.to(iids.device)]
+    topk5_iids = iids[topk5.indices.to(iids.device)]
+    topk1_iids = iids[topk1.indices.to(iids.device)]
+
+    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)
+
+
+@torch.no_grad()
+def compute_irtr_recall_with_rerank(pl_module, split="test"):
+    world_size = dist.get_world_size()
+    rank = dist.get_rank()
+    
+    if split == "val":
+        rank_zero_info("Use val set...")
+        text_dset = pl_module.trainer.datamodule.dms[0].make_no_false_val_dset()
+    else:
+        rank_zero_info("Use test set...")
+        text_dset = pl_module.trainer.datamodule.dms[0].make_no_false_test_dset()
+
+    text_dset.tokenizer = pl_module.trainer.datamodule.dms[0].tokenizer
+    text_loader = torch.utils.data.DataLoader(
+        text_dset,
+        batch_size=32,
+        num_workers=2, #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,
+        ),
+    )
+
+    if split == "val":
+        image_dset = pl_module.trainer.datamodule.dms[0].make_no_false_val_dset(
+            image_only=True
+        )
+    else:
+        image_dset = pl_module.trainer.datamodule.dms[0].make_no_false_test_dset(
+            image_only=True
+        )
+
+    image_dset.tokenizer = pl_module.trainer.datamodule.dms[0].tokenizer
+    image_loader = torch.utils.data.DataLoader(
+        image_dset,
+        batch_size=32,
+        num_workers=2, #pl_module.hparams.config["num_workers"],
+        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()
+    text_ids_list = list()
+    text_masks_list = list()
+    text_labels_list = list()
+    for pre in text_preload:
+        tiids += pre["img_index"]
+        text_ids_list.append(pre["text_ids"])
+        text_masks_list.append(pre["text_masks"])
+        text_labels_list.append(pre["text_labels"])
+    tiids = torch.tensor(tiids)
+    all_text_ids = torch.cat(text_ids_list)
+    all_text_masks = torch.cat(text_masks_list)
+    all_text_labels = torch.cat(text_labels_list)
+
+    image_preload = list()
+    for _b in tqdm.tqdm(image_loader, desc="image prefetch loop"):
+        image_preload.append(
+            {
+                "image": [_b["image"][0].to(pl_module.device)],
+                "img_index": _b["img_index"],
+            }
+        )
+    iids = list()
+    image_list = list()
+    for pre in image_preload:
+        iids += pre["img_index"]
+        image_list.append(pre["image"][0])
+    iids = torch.tensor(iids)
+    all_image = torch.cat(image_list)
+
+    txt_cls_feats = list()
+    for txt_batch in text_preload:
+        with torch.cuda.amp.autocast():
+            cls_feats = pl_module.infer_text_ft(
+                    {
+                        "text_ids": txt_batch["text_ids"],
+                        "text_masks": txt_batch["text_masks"],
+                        "text_labels": txt_batch["text_labels"],
+                    }
+                )["cls_feats"]
+        txt_cls_feats.append(cls_feats)
+
+    img_cls_feats = list()
+    for img_batch in image_preload:
+        with torch.cuda.amp.autocast():
+            cls_feats = pl_module.infer_image_ft(
+                    {
+                        "image": img_batch["image"],
+                    }
+                )["cls_feats"]
+        img_cls_feats.append(cls_feats)
+
+    txt_cls_feats = torch.cat(txt_cls_feats)
+    img_cls_feats = torch.cat(img_cls_feats)
+
+    scores = img_cls_feats @ txt_cls_feats.t()
+
+    rank_zero_info("scores.size(): {}".format(scores.size(), split))
+
+
+    scores_i2t = torch.full((len(iids), len(tiids)), -100.0).to(pl_module.device)
+    
+    k_test = pl_module.hparams.config["k_test"]
+    num_tasks = world_size
+    step = scores.size(0) // num_tasks + 1
+    start = rank * step
+    end = min(scores.size(0), start+step)
+
+    for i, sims in enumerate(scores[start:end]): 
+        if i%100 == 0:
+            rank_zero_info("TR Rerank: {}".format(i))
+        topk_sim, topk_idx = sims.topk(k=k_test, dim=0)
+        cur_images = all_image[start+i].repeat(k_test, 1, 1, 1)
+        cur_text_ids = all_text_ids[topk_idx]
+        cur_text_masks = all_text_masks[topk_idx]
+        cur_text_labels = all_text_labels[topk_idx]
+        
+        cur_rerank_batch = {"image":[cur_images], "text_ids":cur_text_ids, "text_labels":cur_text_labels, "text_masks":cur_text_masks}
+        infer_rerank = pl_module.infer(cur_rerank_batch, mask_text=False, mask_image=False)
+        itm_logits = pl_module.itm_score(infer_rerank["cls_feats"])
+        itm_scores = itm_logits[:,1]
+        scores_i2t[start+i,topk_idx] = itm_scores
+
+    scores = scores.t()
+    scores_t2i = torch.full((len(tiids), len(iids)), -100.0).to(pl_module.device)
+    
+    step = scores.size(0) // num_tasks + 1
+    start = rank * step
+    end = min(scores.size(0), start+step)    
+    
+    for i,sims in enumerate(scores[start:end]): 
+        topk_sim, topk_idx = sims.topk(k=k_test, dim=0)
+        cur_images = all_image[topk_idx]
+        cur_text_ids = all_text_ids[start+i].repeat(k_test, 1) 
+        cur_text_masks = all_text_masks[start+i].repeat(k_test, 1) 
+        cur_text_labels = all_text_labels[start+i].repeat(k_test, 1)
+        cur_rerank_batch = {"image":[cur_images], "text_ids":cur_text_ids, "text_labels":cur_text_labels, "text_masks":cur_text_masks}
+        infer_rerank = pl_module.infer(cur_rerank_batch, mask_text=False, mask_image=False)
+        itm_logits = pl_module.itm_score(infer_rerank["cls_feats"])
+        itm_scores = itm_logits[:,1]
+        scores_t2i[start+i, topk_idx] = itm_scores
+
+    dist.barrier()   
+    torch.distributed.all_reduce(scores_i2t, op=torch.distributed.ReduceOp.SUM) 
+    torch.distributed.all_reduce(scores_t2i, op=torch.distributed.ReduceOp.SUM)    
+
+
+    scores_t2i = scores_t2i + scores
+    scores_i2t = scores_i2t + scores.t()
+
+    topk10 = scores_i2t.topk(10, dim=1)
+    topk5 = scores_i2t.topk(5, dim=1)
+    topk1 = scores_i2t.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_t2i.topk(10, dim=1)
+    topk5 = scores_t2i.topk(5, dim=1)
+    topk1 = scores_t2i.topk(1, dim=1)
+    topk10_iids = iids[topk10.indices]
+    topk5_iids = iids[topk5.indices]
+    topk1_iids = iids[topk1.indices]
+
+    ir_r10 = (tiids.unsqueeze(1) == topk10_iids).float().max(dim=1)[0].mean()
+    ir_r5 = (tiids.unsqueeze(1) == topk5_iids).float().max(dim=1)[0].mean()
+    ir_r1 = (tiids.unsqueeze(1) == topk1_iids).float().max(dim=1)[0].mean()
+
+    return (ir_r1, ir_r5, ir_r10, tr_r1, tr_r5, tr_r10)
+
+
+def init_weights(module):
+    if isinstance(module, (nn.Linear, nn.Embedding)):
+        module.weight.data.normal_(mean=0.0, std=0.02)
+    elif isinstance(module, nn.LayerNorm):
+        module.bias.data.zero_()
+        module.weight.data.fill_(1.0)
+
+    if isinstance(module, nn.Linear) and module.bias is not None:
+        module.bias.data.zero_()
+
+
+def vqa_test_step(pl_module, batch, output):
+    id2answer = (
+        pl_module.trainer.datamodule.dm_dicts["vqa_trainval"].id2answer
+        if "vqa_trainval" in pl_module.trainer.datamodule.dm_dicts
+        else pl_module.trainer.datamodule.dm_dicts["vqa"].id2answer
+    )
+    vqa_logits = output["vqa_logits"]
+    vqa_preds = vqa_logits.argmax(dim=-1)
+    vqa_preds = [id2answer[pred.item()] for pred in vqa_preds]
+    questions = batch["text"]
+    qids = batch["qid"]
+    return {"qids": qids, "preds": vqa_preds}
+
+
+def vqa_test_wrapup(outs, model_name, output_dir):
+    rank = torch.distributed.get_rank()
+    qids, preds = list(), list()
+    for out in outs:
+        qids += out["qids"]
+        preds += out["preds"]
+
+    rets = list()
+    for qid, pred in zip(qids, preds):
+        rets.append({"question_id": qid, "answer": pred})
+    with open(f"vqa_submit_{rank}.json", "w") as fp:
+        json.dump(rets, fp, indent=4)
+
+    torch.distributed.barrier()
+
+    if rank == 0:
+        jsons = list()
+        paths = list(glob.glob("vqa_submit_*.json"))
+        for path in paths:
+            with open(path, "r") as fp:
+                jsons += json.load(fp)
+        # os.makedirs("result", exist_ok=True)
+        os.makedirs(output_dir, exist_ok=True)
+        with open(f"{output_dir}/vqa_submit_{model_name}.json", "w") as fp:
+            json.dump(jsons, fp, indent=4)
+
+    torch.distributed.barrier()
+    os.remove(f"vqa_submit_{rank}.json")