mirror of
https://github.com/brycedrennan/imaginAIry
synced 2024-10-31 03:20:40 +00:00
119 lines
4.0 KiB
Python
119 lines
4.0 KiB
Python
"""Classes for visual-language reasoning"""
|
|
|
|
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
|