import torch.nn as nn
import torch
from transformers import T5ForConditionalGeneration, ViTModel
import pytorch_lightning as pl
# Defining the pytorch model
class LaTr_for_pretraining(nn.Module):
def __init__(self, config, classify=False):
super(LaTr_for_pretraining, self).__init__()
self.vocab_size = config['vocab_size']
model = T5ForConditionalGeneration.from_pretrained(config['t5_model'])
# Removing the Embedding layer
dummy_encoder = list(nn.Sequential(
*list(model.encoder.children())[1:]).children())
# Removing the Embedding Layer
dummy_decoder = list(nn.Sequential(
*list(model.decoder.children())[1:]).children())
# Using the T5 Encoder
self.list_encoder = nn.Sequential(*list(dummy_encoder[0]))
self.residue_encoder = nn.Sequential(*list(dummy_encoder[1:]))
self.list_decoder = nn.Sequential(*list(dummy_decoder[0]))
self.residue_decoder = nn.Sequential(*list(dummy_decoder[1:]))
# We use the embeddings of T5 for encoding the tokenized words
self.language_emb = nn.Embedding.from_pretrained(model.shared.weight)
self.top_left_x = nn.Embedding(
config['max_2d_position_embeddings'], config['hidden_state'])
self.bottom_right_x = nn.Embedding(
config['max_2d_position_embeddings'], config['hidden_state'])
self.top_left_y = nn.Embedding(
config['max_2d_position_embeddings'], config['hidden_state'])
self.bottom_right_y = nn.Embedding(
config['max_2d_position_embeddings'], config['hidden_state'])
self.width_emb = nn.Embedding(
config['max_2d_position_embeddings'], config['hidden_state'])
self.height_emb = nn.Embedding(
config['max_2d_position_embeddings'], config['hidden_state'])
self.classify = classify
self.classification_layer = nn.Linear(
config['hidden_state'], config['classes'])
def forward(self, tokens, coordinates, predict_proba=False, predict_class=False):
batch_size = len(tokens)
embeded_feature = self.language_emb(tokens)
top_left_x_feat = self.top_left_x(coordinates[:, :, 0])
top_left_y_feat = self.top_left_y(coordinates[:, :, 1])
bottom_right_x_feat = self.bottom_right_x(coordinates[:, :, 2])
bottom_right_y_feat = self.bottom_right_y(coordinates[:, :, 3])
width_feat = self.width_emb(coordinates[:, :, 4])
height_feat = self.height_emb(coordinates[:, :, 5])
total_feat = embeded_feature + top_left_x_feat + top_left_y_feat + \
bottom_right_x_feat + bottom_right_y_feat + width_feat + height_feat
# Extracting the feature
for layer in self.list_encoder:
total_feat = layer(total_feat)[0]
total_feat = self.residue_encoder(total_feat)
for layer in self.list_decoder:
total_feat = layer(total_feat)[0]
total_feat = self.residue_decoder(total_feat)
if self.classify:
total_feat = self.classification_layer(total_feat)
if predict_proba:
return total_feat.softmax(axis=-1)
if predict_class:
return total_feat.argmax(axis=-1)
return total_feat
class LaTr_for_finetuning(nn.Module):
def __init__(self, config, address_to_pre_trained_weights=None):
super(LaTr_for_finetuning, self).__init__()
self.config = config
self.vocab_size = config['vocab_size']
self.pre_training_model = LaTr_for_pretraining(config)
if address_to_pre_trained_weights is not None:
self.pre_training_model.load_state_dict(
torch.load(address_to_pre_trained_weights))
self.vit = ViTModel.from_pretrained(
"google/vit-base-patch16-224-in21k")
# In the fine-tuning stage of vit, except the last layer, all the layers were freezed
self.classification_head = nn.Linear(
config['hidden_state'], config['classes'])
def forward(self, lang_vect, spatial_vect, quest_vect, img_vect):
# The below block of code calculates the language and spatial featuer
embeded_feature = self.pre_training_model.language_emb(lang_vect)
top_left_x_feat = self.pre_training_model.top_left_x(
spatial_vect[:, :, 0])
top_left_y_feat = self.pre_training_model.top_left_y(
spatial_vect[:, :, 1])
bottom_right_x_feat = self.pre_training_model.bottom_right_x(
spatial_vect[:, :, 2])
bottom_right_y_feat = self.pre_training_model.bottom_right_y(
spatial_vect[:, :, 3])
width_feat = self.pre_training_model.width_emb(spatial_vect[:, :, 4])
height_feat = self.pre_training_model.height_emb(spatial_vect[:, :, 5])
spatial_lang_feat = embeded_feature + top_left_x_feat + top_left_y_feat + \
bottom_right_x_feat + bottom_right_y_feat + width_feat + height_feat
# Extracting the image feature, using the Vision Transformer
img_feat = self.vit(img_vect).last_hidden_state
# Extracting the question vector
quest_feat = self.pre_training_model.language_emb(quest_vect)
# Concating the three features, and then passing it through the T5 Transformer
final_feat = torch.cat(
[img_feat, spatial_lang_feat, quest_feat], axis=-2)
# Passing through the T5 Transformer
for layer in self.pre_training_model.list_encoder:
final_feat = layer(final_feat)[0]
final_feat = self.pre_training_model.residue_encoder(final_feat)
for layer in self.pre_training_model.list_decoder:
final_feat = layer(final_feat)[0]
final_feat = self.pre_training_model.residue_decoder(final_feat)
answer_vector = self.classification_head(
final_feat)[:, :self.config['seq_len'], :]
return answer_vector
def polynomial(base_lr, iter, max_iter=1e5, power=1):
return base_lr * ((1 - float(iter) / max_iter) ** power)
class LaTrForVQA(pl.LightningModule):
def __init__(self, config, learning_rate=1e-4, max_steps=100000//2):
super(LaTrForVQA, self).__init__()
self.config = config
self.save_hyperparameters()
self.latr = LaTr_for_finetuning(config)
self.training_losses = []
self.validation_losses = []
self.max_steps = max_steps
def configure_optimizers(self):
return torch.optim.AdamW(self.parameters(), lr=self.hparams['learning_rate'])
def forward(self, batch_dict):
boxes = batch_dict['boxes']
img = batch_dict['img']
question = batch_dict['question']
words = batch_dict['tokenized_words']
answer_vector = self.latr(lang_vect=words,
spatial_vect=boxes,
img_vect=img,
quest_vect=question
)
return answer_vector
def calculate_metrics(self, prediction, labels):
# Calculate the accuracy score between the prediction and ground label for a batch, with considering the pad sequence
batch_size = len(prediction)
ac_score = 0
for (pred, gt) in zip(prediction, labels):
ac_score += calculate_acc_score(pred.detach().cpu(),
gt.detach().cpu())
ac_score = ac_score/batch_size
return ac_score
def training_step(self, batch, batch_idx):
answer_vector = self.forward(batch)
# https://discuss.huggingface.co/t/bertformaskedlm-s-loss-and-scores-how-the-loss-is-computed/607/2
loss = nn.CrossEntropyLoss(ignore_index=0)(
answer_vector.reshape(-1, self.config['classes']), batch['answer'].reshape(-1))
_, preds = torch.max(answer_vector, dim=-1)
# Calculating the accuracy score
train_acc = self.calculate_metrics(preds, batch['answer'])
train_acc = torch.tensor(train_acc)
# Logging
self.log('train_ce_loss', loss, prog_bar=True)
self.log('train_acc', train_acc, prog_bar=True)
self.training_losses.append(loss.item())
return loss
def validation_step(self, batch, batch_idx):
logits = self.forward(batch)
loss = nn.CrossEntropyLoss(ignore_index=0)(
logits.reshape(-1, self.config['classes']), batch['answer'].reshape(-1))
_, preds = torch.max(logits, dim=-1)
# Validation Accuracy
val_acc = self.calculate_metrics(preds.cpu(), batch['answer'].cpu())
val_acc = torch.tensor(val_acc)
# Logging
self.log('val_ce_loss', loss, prog_bar=True)
self.log('val_acc', val_acc, prog_bar=True)
self.validation_losses.append(loss.item())
return {'val_loss': loss, 'val_acc': val_acc}
def optimizer_step(self, epoch_nb, batch_nb, optimizer, optimizer_i, opt_closure=None, on_tpu=False,
using_native_amp=False, using_lbfgs=False):
# Warmup for 1000 steps
if self.trainer.global_step < 1000:
lr_scale = min(1., float(self.trainer.global_step + 1) / 1000.)
for pg in optimizer.param_groups:
pg['lr'] = lr_scale * self.hparams.learning_rate
# Linear Decay
else:
for pg in optimizer.param_groups:
pg['lr'] = polynomial(
self.hparams.learning_rate, self.trainer.global_step, max_iter=self.max_steps)
optimizer.step(opt_closure)
optimizer.zero_grad()
def validation_epoch_end(self, outputs):
val_loss = torch.stack([x['val_loss'] for x in outputs]).mean()
val_acc = torch.stack([x['val_acc'] for x in outputs]).mean()
self.log('val_loss_epoch_end', val_loss, on_epoch=True, sync_dist=True)
self.log('val_acc_epoch_end', val_acc, on_epoch=True, sync_dist=True)