svitt/model.py

from svitt.utils import (
    interpolate_pos_embed,
    interpolate_pos_relative_bias_beit_3d,
)
from omegaconf import OmegaConf
from transformers import ViTModel, ViTConfig
from svitt.sparse_config import BertConfig, BeitConfig
from svitt.sparse_xbeit import BeitModel
from svitt.sparse_xbert import BertModel, BertForMaskedLM

import torch
from torch import nn
import torch.nn.functional as F


class SViTT(nn.Module):
    """Common utils shared by pretraining and downstream retrieval"""
    def __init__(self, config=None, tokenizer=None, pretrain=True, **kwargs):
        super().__init__()
        self.config = config
        self.tokenizer = tokenizer
        self.embed_dim = config.embed_dim
        self.vision_width = 768
        self.text_width = 768
        self.pretrain = pretrain

        self.vision_encoder, self.vision_layernorm = self.build_vision_encoder()
        self.text_encoder = self.build_text_encoder()

        self.vision_proj = nn.Linear(self.vision_width, self.embed_dim)
        self.text_proj = nn.Linear(self.text_width, self.embed_dim)

        self.temp = nn.Parameter(torch.ones([]) * config.temp)
        self.itm_head = nn.Linear(self.text_width, 2)


    def build_text_encoder(self):

        bert_config = BertConfig.from_json_file(self.config.bert_config)

        # Override params for sparse vision encoder
        model_args = getattr(self.config, 'text_encoder_args', {})
        if model_args:
            model_args = OmegaConf.to_object(model_args)
            bert_config.update(model_args)

        if self.pretrain:
            text_encoder, _ = BertForMaskedLM.from_pretrained(
                self.config.text_encoder, config=bert_config,
                output_loading_info=True
            )
        else:
            text_encoder, _ = BertModel.from_pretrained(
                self.config.text_encoder, config=bert_config,
                add_pooling_layer=False, output_loading_info=True
            )
        return text_encoder

    def build_vision_encoder(self):
        # if self.config.vit_type in ["beit", "deit", "vit", "vit32"]:
        if self.config.vit_type in ["beit"]:
            vision_encoder = self.build_huggingface_vit_with_image_size(
                self.config.vit_name_or_pretrained_path, self.config.image_res,)
        else:
            raise ValueError(f"Unknown vit type {self.config.vit_type}")

        # add layernorm for normalizing BEiT outputs hidden states
        vision_layernorm = None
        if self.config.vit_type == "beit":
            vision_layernorm = nn.LayerNorm(self.vision_width, eps=1e-12)
        return vision_encoder, vision_layernorm

    # @classmethod
    # def build_huggingface_vit_with_image_size(cls, model_card: str, image_size: int):
    def build_huggingface_vit_with_image_size(self, model_card: str, image_size: int):
        """Build a vit model from huggingface hub, also interpolate pos_embed when needed.

        Args:
            model_card: name in huggingface hub, e.g., `facebook/deit-base-patch16-224`
            image_size: new image size, may be different from pre-training image_size of `model_card`

        ref: https://github.com/huggingface/transformers/issues/12167#issuecomment-861356232
        """
        is_beit = "beit" in model_card
        if "beit" in model_card:
            model_cls, config_cls = BeitModel, BeitConfig
        elif "deit" in model_card or "vit" in model_card:
            # the deit model we use is loaded in vit arch,
            # see https://huggingface.co/facebook/deit-base-patch16-224#how-to-use
            model_cls, config_cls = ViTModel, ViTConfig
        else:
            raise ValueError(f"Unexpected model_card: {model_card}")

        # BEiT uses average pooled tokens instead of [CLS] used by other models
        tmp_model = model_cls.from_pretrained(model_card, add_pooling_layer=is_beit)
        state_dict = tmp_model.state_dict()
        del tmp_model

        # Override params for sparse vision encoder
        model_args = getattr(self.config, 'vision_encoder_args', {})
        if model_args:
            model_args = OmegaConf.to_object(model_args)
        model_config = config_cls.from_pretrained(
            model_card, 
            image_size=image_size, 
            **model_args,
        )
        model = model_cls(config=model_config, add_pooling_layer=is_beit, num_frames=self.config.video_input.num_frames)
        if is_beit:
            # interpolate relative pos bias
            state_dict = interpolate_pos_relative_bias_beit_3d(
                state_dict_old=state_dict,
                state_dict_new=model.state_dict(),
                patch_shape_new=model.window_size
            )
        else:
            # interpolate pos_embed and load weights to new model
            state_dict["embeddings.position_embeddings"] = interpolate_pos_embed(
                pos_embed_old=state_dict["embeddings.position_embeddings"],
                pos_embed_new=model.embeddings.position_embeddings,
                num_patches_new=model.embeddings.patch_embeddings.num_patches
            )
        msg = model.load_state_dict(state_dict, strict=False)
        return model

    def get_text_encoder(self):
        """get text encoder, used for text and cross-modal encoding"""
        encoder = self.text_encoder
        return encoder.bert if hasattr(encoder, "bert") else encoder

    def encode_image(self, video, output_token_idx=False, output_attentions=False):
        video_embeds = self.vision_encoder(video, output_token_idx=output_token_idx, output_attentions=output_attentions)  # (bsz, seq_len, d)
        if self.vision_layernorm is not None:  # only for BEiT mean-pooling
            video_embeds.last_hidden_state = self.vision_layernorm(video_embeds.last_hidden_state)
        if output_token_idx:
            token_idx = video_embeds.token_idx
        
        if output_attentions:
            attentions = video_embeds.attentions
            
        if self.config.vit_type == "beit":
            pooled_video_embeds = video_embeds.pooler_output  # (bsz*num_frms, d)
            video_embeds = video_embeds.last_hidden_state  # (bsz*num_frms, L, d)
        else:
            video_embeds = video_embeds.last_hidden_state
            pooled_video_embeds = video_embeds[:, 0]
        
        outputs = (video_embeds, pooled_video_embeds)

        if output_token_idx:
            outputs += (token_idx,)
        
        if output_attentions:
            outputs += (attentions,)

        return outputs

    def _encode_image(self, image):
        bsz, num_frms, c, h, w = image.shape  # `num_frms` could be changing for image (=1) or video (e.g., =4)
        image = image.view(bsz*num_frms, c, h, w)
        image_embeds = self.vision_encoder(image)
        if self.vision_layernorm is not None:  # only for BEiT mean-pooling
            image_embeds.last_hidden_state = self.vision_layernorm(image_embeds.last_hidden_state)

        if self.config.vit_type == "beit":
            pooled_image_embeds = image_embeds.pooler_output  # (bsz*num_frms, d)
            image_embeds = image_embeds.last_hidden_state  # (bsz*num_frms, L, d)
        else:
            image_embeds = image_embeds.last_hidden_state
            pooled_image_embeds = image_embeds[:, 0]

        image_embeds = image_embeds.view(bsz, num_frms, -1, self.vision_width)  # (bsz, num_frms, L, d)
        pooled_image_embeds = pooled_image_embeds.view(bsz, num_frms, self.vision_width) \
            if pooled_image_embeds is not None else None  # (bsz, num_frms, d)
        return image_embeds, pooled_image_embeds

    def encode_text(self, text):
        text_output = self.get_text_encoder()(
            text.input_ids,
            attention_mask=text.attention_mask,
            return_dict=True,
            mode='text'
        )
        text_embeds = text_output.last_hidden_state 
        pooled_text_embeds = text_embeds[:, 0] 
        return text_embeds, pooled_text_embeds

    @torch.no_grad()
    def clip_contrastive_temperature(self, min_val=0.001, max_val=0.5):
        """Seems only used during pre-training"""
        self.temp.clamp_(min_val, max_val)

    @torch.no_grad()
    def get_mask(self, sim, idx=None, normalize=False):
        """
        sim: (N, N)
        idx: (N, )
        normalize: bool, make row sum equal to 1
        """
        if idx is not None:
            idx = idx.view(-1, 1)
            mask = torch.eq(idx, idx.T).to(sim.dtype)
            if normalize:  
                mask = mask / mask.sum(1, keepdim=True)
        else:
            mask = torch.zeros_like(sim)
            mask.fill_diagonal_(1)
        return mask  # `1` mark valid/matched location

    def get_contrastive_loss(self, pooled_image_embeds, pooled_text_embeds, idx=None):
        sim_i2t, sim_t2i = self.get_sim(
            pooled_image_embeds, pooled_text_embeds, t=self.temp)

        with torch.no_grad():
            sim_i2t_targets = self.get_mask(sim_i2t, idx=idx, normalize=True)
            sim_t2i_targets = sim_i2t_targets

        loss_i2t = -torch.sum(
            F.log_softmax(sim_i2t, dim=1) * sim_i2t_targets, dim=1).mean()
        loss_t2i = -torch.sum(
            F.log_softmax(sim_t2i, dim=1) * sim_t2i_targets, dim=1).mean()

        loss_ita = (loss_i2t + loss_t2i) / 2
        return loss_ita, sim_i2t, sim_t2i

    def get_sim(self, pooled_image_embeds, pooled_text_embeds, t=1):
        """
        Args:
            pooled_image_embeds: (bsz, num_frms, d)
            pooled_text_embeds: (bsz, d)
            t: temperature
        """
        image_proj = self.vision_proj
        text_proj = self.text_proj

        image_feat = F.normalize(image_proj(pooled_image_embeds), dim=-1)
        text_feat = F.normalize(text_proj(pooled_text_embeds), dim=-1)

        if image_feat.ndim == 3:
            sim_i2t = torch.einsum("mld,nd->mln", image_feat, text_feat).mean(1) / t  # (N, N)
        else:
            sim_i2t = torch.einsum("md,nd ->mn", image_feat, text_feat) / t  # (N, N)
        sim_t2i = sim_i2t.T
        return sim_i2t, sim_t2i

    def get_itm_loss(self, 
            sim_i2t, 
            sim_t2i, 
            text_embeds, 
            text_atts, 
            image_embeds, 
            image_atts, 
            idx=None,
        ):
        """
        sim_i2t, sim_t2i: (N, N)
        text_embeds, text_atts, image_embeds, image_atts: (N, *)
        idx: (N, )
        """
        bsz = len(sim_i2t)

        with torch.no_grad():
            weights_i2t = F.softmax(sim_i2t+1e-4, dim=1)  # (N, N)
            weights_t2i = F.softmax(sim_t2i+1e-4, dim=1)

            mask = self.get_mask(sim_i2t, idx=idx).bool()
            weights_i2t.masked_fill_(mask, 0)
            weights_t2i.masked_fill_(mask, 0)

        # select a negative image for each text
        if self.config.itm_hard_neg:
            img_neg_indices = torch.multinomial(weights_t2i, 1).squeeze() #RuntimeError: invalid multinomial distribution (sum of probabilities <= 0) 
        else:
            img_neg_indices = self.get_rand_indices(mask, 1).squeeze()            

        image_embeds_neg = image_embeds[img_neg_indices]

        # select a negative text for each image
        if self.config.itm_hard_neg:
            txt_neg_indices = torch.multinomial(weights_i2t, 1).squeeze()
        else:
            txt_neg_indices = self.get_rand_indices(mask, 1).squeeze()

        text_embeds_neg = text_embeds[txt_neg_indices]
        text_atts_neg = text_atts[txt_neg_indices]  # (N, L, d)

        # embedding on local gpu
        _text_embeds = text_embeds  
        _text_atts = text_atts
        _image_embeds = image_embeds
        _image_atts = image_atts
        # concat embeddings
        text_embeds_all = torch.cat([_text_embeds, _text_embeds, text_embeds_neg], dim=0)
        text_atts_all = torch.cat([_text_atts, _text_atts, text_atts_neg], dim=0)
        image_embeds_all = torch.cat([_image_embeds, image_embeds_neg, _image_embeds], dim=0)
        image_atts_all = torch.cat([_image_atts, _image_atts, _image_atts], dim=0)

        text_encoder = self.get_text_encoder()
        output = text_encoder(
            encoder_embeds=text_embeds_all,
            attention_mask=text_atts_all,
            encoder_hidden_states=image_embeds_all,
            encoder_attention_mask=image_atts_all,
            return_dict=True,
            mode='fusion',
        )

        itm_embeds = output.last_hidden_state[:, 0]  # pos (N, d) + neg (2N, d)

        loss_itm = self._get_itm_loss(itm_embeds, enc=self.itm_head)
        itm_embeds_pos = itm_embeds[:bsz]  # (N, d)

        return loss_itm, itm_embeds_pos

    def _get_itm_loss(self, itm_embeds, enc):
        """
        itm_embeds: (3*N, D)
        enc: nn.Module that projects cls_embeds
        """
        itm_scores = enc(itm_embeds)  # (3*N, 2)
        bs = itm_scores.size(0) // 3
        itm_labels = itm_scores.new_ones(3*bs, dtype=torch.long)
        itm_labels[bs:] = 0
        loss_itm = F.cross_entropy(itm_scores, itm_labels)
        return loss_itm

    def get_rand_indices(self, mask, k):
        """
        Args:
            mask: (N, L) 0 indicates the positions that we can sample, 1 otherwise
            k: #indices to sample at each row
        Returns:
            (N, k) indices
        """
        mask = mask.float()
        mask = mask - 10000 * mask
        mask += torch.randn_like(mask)
        _, indices = torch.sort(mask, dim=1, descending=True)
        indices = indices[:, :k].contiguous()  
        return indices