Upload modeling.py

modeling.py

@@ -0,0 +1,251 @@
+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)