imaginAIry/imaginairy/vendored/blip/blip_nlvr.py

import torch
import torch.nn.functional as F
from models.blip import create_vit, init_tokenizer, is_url
from models.med import BertConfig
from models.nlvr_encoder import BertModel
from models.vit import interpolate_pos_embed
from timm.models.hub import download_cached_file
from torch import nn


class BLIP_NLVR(nn.Module):
    def __init__(
        self,
        med_config="configs/med_config.json",
        image_size=480,
        vit="base",
        vit_grad_ckpt=False,
        vit_ckpt_layer=0,
    ):
        """
        Args:
            med_config (str): path for the mixture of encoder-decoder model's configuration file
            image_size (int): input image size
            vit (str): model size of vision transformer
        """
        super().__init__()

        self.visual_encoder, vision_width = create_vit(
            vit, image_size, vit_grad_ckpt, vit_ckpt_layer, drop_path_rate=0.1
        )
        self.tokenizer = init_tokenizer()
        med_config = BertConfig.from_json_file(med_config)
        med_config.encoder_width = vision_width
        self.text_encoder = BertModel(config=med_config, add_pooling_layer=False)

        self.cls_head = nn.Sequential(
            nn.Linear(
                self.text_encoder.config.hidden_size,
                self.text_encoder.config.hidden_size,
            ),
            nn.ReLU(),
            nn.Linear(self.text_encoder.config.hidden_size, 2),
        )

    def forward(self, image, text, targets, train=True):

        image_embeds = self.visual_encoder(image)
        image_atts = torch.ones(image_embeds.size()[:-1], dtype=torch.long).to(
            image.device
        )
        image0_embeds, image1_embeds = torch.split(image_embeds, targets.size(0))

        text = self.tokenizer(text, padding="longest", return_tensors="pt").to(
            image.device
        )
        text.input_ids[:, 0] = self.tokenizer.enc_token_id

        output = self.text_encoder(
            text.input_ids,
            attention_mask=text.attention_mask,
            encoder_hidden_states=[image0_embeds, image1_embeds],
            encoder_attention_mask=[
                image_atts[: image0_embeds.size(0)],
                image_atts[image0_embeds.size(0) :],
            ],
            return_dict=True,
        )
        hidden_state = output.last_hidden_state[:, 0, :]
        prediction = self.cls_head(hidden_state)

        if train:
            loss = F.cross_entropy(prediction, targets)
            return loss
        else:
            return prediction


def blip_nlvr(pretrained="", **kwargs):
    model = BLIP_NLVR(**kwargs)
    if pretrained:
        model, msg = load_checkpoint(model, pretrained)
        print("missing keys:")
        print(msg.missing_keys)
    return model


def load_checkpoint(model, url_or_filename):
    if is_url(url_or_filename):
        cached_file = download_cached_file(
            url_or_filename, check_hash=False, progress=True
        )
        checkpoint = torch.load(cached_file, map_location="cpu")
    elif os.path.isfile(url_or_filename):
        checkpoint = torch.load(url_or_filename, map_location="cpu")
    else:
        raise RuntimeError("checkpoint url or path is invalid")
    state_dict = checkpoint["model"]

    state_dict["visual_encoder.pos_embed"] = interpolate_pos_embed(
        state_dict["visual_encoder.pos_embed"], model.visual_encoder
    )

    for key in list(state_dict.keys()):
        if "crossattention.self." in key:
            new_key0 = key.replace("self", "self0")
            new_key1 = key.replace("self", "self1")
            state_dict[new_key0] = state_dict[key]
            state_dict[new_key1] = state_dict[key]
        elif "crossattention.output.dense." in key:
            new_key0 = key.replace("dense", "dense0")
            new_key1 = key.replace("dense", "dense1")
            state_dict[new_key0] = state_dict[key]
            state_dict[new_key1] = state_dict[key]

    msg = model.load_state_dict(state_dict, strict=False)
    print("load checkpoint from %s" % url_or_filename)
    return model, msg