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)