fairseq/examples/MMPT/mmpt/modules/vectorpool.py

# Copyright (c) Facebook, Inc. All Rights Reserved

import torch
import os
import numpy as np
import pickle

from . import retri
from ..utils import get_local_rank


class VectorPool(object):
    """
    Base class of retrieval space.
    """

    def __init__(self, config):
        from transformers import AutoConfig
        self.hidden_size = AutoConfig.from_pretrained(
            config.dataset.bert_name).hidden_size
        self.retriever_cls = getattr(retri, config.retriever_cls)

    def __call__(self, sample, **kwargs):
        raise NotImplementedError

    def build_retriver(
        self,
        retriever_cls=None,
        hidden_size=None,
        centroids=512,
        db_type="flatl2",
        examples_per_cent_to_train=48
    ):

        """merge results from multiple gpus and return a retriver.."""
        self.retriver = retriever_cls(
            hidden_size, centroids, db_type, examples_per_cent_to_train)
        return self.retriver

    def __repr__(self):
        if hasattr(self, "retriver"):
            retriver_name = str(len(self.retriver))
        else:
            retriver_name = "no retriver field yet"
        return self.__class__.__name__ \
            + "(" + retriver_name + ")"


class VideoVectorPool(VectorPool):
    """
    average clips of a video as video representation.
    """
    def __init__(self, config):
        super().__init__(config)
        self.build_retriver(self.retriever_cls, self.hidden_size)

    def __call__(self, sample, subsampling, **kwargs):
        hidden_states = (
            sample["pooled_video"] + sample["pooled_text"]) / 2.
        hidden_states = hidden_states.view(
            -1, subsampling,
            hidden_states.size(-1))
        hidden_states = torch.mean(hidden_states, dim=1)
        hidden_states = hidden_states.cpu().detach().numpy()
        video_ids = []
        for offset_idx, video_id in enumerate(sample["video_id"]):
            if isinstance(video_id, tuple) and len(video_id) == 3:
                # a sharded video_id.
                video_id = video_id[0]
            video_ids.append(video_id)
        assert len(video_ids) == len(hidden_states)
        self.retriver.add(
            hidden_states.astype("float32"),
            video_ids
        )


class DistributedVectorPool(VectorPool):
    """
    support sync of multiple gpus/nodes.
    """
    def __init__(self, config):
        super().__init__(config)
        self.out_dir = os.path.join(
            config.fairseq.checkpoint.save_dir,
            "retri")
        os.makedirs(self.out_dir, exist_ok=True)
        self.hidden_states = []
        self.video_ids = []

    def build_retriver(
        self,
        retriever_cls=None,
        hidden_size=None,
        centroids=4096,
        db_type="flatl2",
        examples_per_cent_to_train=48
    ):
        if retriever_cls is None:
            retriever_cls = self.retriever_cls
        if hidden_size is None:
            hidden_size = self.hidden_size
        """merge results from multiple gpus and return a retriver.."""
        if torch.distributed.is_initialized():
            self.save()
            # sync saving.
            torch.distributed.barrier()
            world_size = torch.distributed.get_world_size()
        else:
            world_size = 1
        self.retriver = retriever_cls(
            hidden_size, centroids, db_type, examples_per_cent_to_train)
        # each gpu process has its own retriever.
        for local_rank in range(world_size):
            if get_local_rank() == 0:
                print("load local_rank", local_rank)
            hidden_states, video_ids = self.load(local_rank)
            hidden_states = hidden_states.astype("float32")
            self.retriver.add(hidden_states, video_ids)
        return self.retriver

    def load(self, local_rank):
        hidden_states = np.load(
            os.path.join(
                self.out_dir,
                "hidden_state" + str(local_rank) + ".npy"
            )
        )

        with open(
            os.path.join(
                self.out_dir, "video_id" + str(local_rank) + ".pkl"),
                "rb") as fr:
            video_ids = pickle.load(fr)
        return hidden_states, video_ids

    def save(self):
        hidden_states = np.vstack(self.hidden_states)
        assert len(hidden_states) == len(self.video_ids), "{}, {}".format(
            len(hidden_states),
            len(self.video_ids)
        )
        local_rank = torch.distributed.get_rank() \
            if torch.distributed.is_initialized() else 0

        np.save(
            os.path.join(
                self.out_dir,
                "hidden_state" + str(local_rank) + ".npy"),
            hidden_states)

        with open(
            os.path.join(
                self.out_dir,
                "video_id" + str(local_rank) + ".pkl"),
                "wb") as fw:
            pickle.dump(
                self.video_ids,
                fw,
                protocol=pickle.HIGHEST_PROTOCOL
            )


class DistributedVideoVectorPool(DistributedVectorPool):
    """
    average clips of a video as video representation.
    """
    def __call__(self, sample, subsampling, **kwargs):
        hidden_states = (
            sample["pooled_video"] + sample["pooled_text"]) / 2.
        hidden_states = hidden_states.view(
            -1, subsampling,
            hidden_states.size(-1))
        hidden_states = torch.mean(hidden_states, dim=1)
        hidden_states = hidden_states.cpu().detach().numpy()
        video_ids = []
        for offset_idx, video_id in enumerate(sample["video_id"]):
            if isinstance(video_id, tuple) and len(video_id) == 3:
                # a sharded video_id.
                video_id = video_id[0]
            video_ids.append(video_id)
        assert len(video_ids) == len(hidden_states)
        self.hidden_states.append(hidden_states)
        self.video_ids.extend(video_ids)


# ------------ the following are deprecated --------------

class TextClipVectorPool(VectorPool):
    def __init__(self, config):
        from transformers import AutoConfig
        hidden_size = AutoConfig.from_pretrained(
            config.dataset.bert_name).hidden_size
        retriever_cls = getattr(retri, config.retriever_cls)
        self.build_retriver(retriever_cls, hidden_size)

    def __call__(self, sample, **kwargs):
        clip_meta = sample["clip_meta"].cpu()
        assert torch.all(torch.le(clip_meta[:, 4], clip_meta[:, 5]))
        text_meta = [tuple(item.tolist()) for item in clip_meta[:, 3:]]

        if hasattr(self, "retriver"):
            # build_retriver is called.
            self.retriver.add(
                sample["pooled_text"].cpu().numpy().astype("float32"),
                text_meta
            )
        else:
            raise NotImplementedError


class MMClipVectorPool(VectorPool):
    """
    Multimodal Clip-level vector pool.
    """
    def __init__(self, out_dir):
        """use hidden_states to store `(video, text)`."""
        """use video_ids to store `(video_id, start, end)`."""
        super().__init__(out_dir)

    def __call__(self, sample, **kwargs):
        pooled_video = sample["pooled_video"].cpu().unsqueeze(1).numpy()
        pooled_text = sample["pooled_text"].cpu().unsqueeze(1).numpy()

        self.hidden_states.append(
            np.concatenate([pooled_video, pooled_text], axis=1)
        )

        video_starts = sample["video_start"].cpu()
        video_ends = sample["video_end"].cpu()
        assert torch.all(torch.le(video_starts, video_ends))

        text_starts = sample["text_start"].cpu()
        text_ends = sample["text_end"].cpu()
        assert torch.all(torch.le(text_starts, text_ends))
        subsample_size = sample["pooled_video"].size(0) // len(sample["video_id"])
        video_ids = [video_id for video_id in sample["video_id"]
                    for _ in range(subsample_size)
        ]
        for video_id, video_start, video_end, text_start, text_end in zip(
                video_ids, video_starts, video_ends, text_starts, text_ends):
            self.video_ids.append((
                video_id,
                (int(video_start), int(video_end)),
                (int(text_start), int(text_end))
            ))