Update ROBERTAmodel.py

ROBERTAmodel.py

@@ -5,9 +5,8 @@ from models import TransformerVisualizer
 from transformers import (
     RobertaForMaskedLM, RobertaForSequenceClassification
 )
-import os
-import time
-
+import os,time
+import torch.autograd.functional as Fgrad
 
 CACHE_DIR  = "/data/hf_cache"
 
@@ -151,7 +150,6 @@ class RoBERTaVisualizer(TransformerVisualizer):
     def get_all_grad_attn_matrix(self, task, sentence, hypothesis='', maskID = None):
         print(task, sentence,  hypothesis)
         print('Tokenize')
-        start = time.time()
         if task == 'mnli':
             inputs = self.tokenizer(sentence, hypothesis, return_tensors='pt', padding=False, truncation=True)
         elif task == 'mlm':
@@ -163,91 +161,96 @@ class RoBERTaVisualizer(TransformerVisualizer):
         tokens = self.tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
         print(tokens)
         inputs = {k: v.to(self.device) for k, v in inputs.items()}
-        print(1,time.time()-start)
-        start = time.time()
 
         print('Input embeddings with grad')
         embedding_layer = self.model.roberta.embeddings.word_embeddings
         inputs_embeds = embedding_layer(inputs["input_ids"])
         inputs_embeds.requires_grad_()
-        
-        print(2,time.time()-start)
-        start = time.time()
+
         print('Forward pass')
         outputs = self.model.roberta(
             inputs_embeds=inputs_embeds,
             attention_mask=inputs["attention_mask"],
             output_attentions=True
         )
+
+
+
+
+
         attentions = outputs.attentions  # list of [1, heads, seq, seq]
-        
-        print(3,time.time()-start)
-        start = time.time()
+
         print('Average attentions per layer')
         mean_attns = [a.squeeze(0).mean(dim=0).detach().cpu() for a in attentions]
-        print(4,time.time()-start)
-        startloop = time.time()
-        start = time.time()
 
         attn_matrices_all = []
         grad_matrices_all = []
         for target_layer in range(len(attentions)):
-            print(5,target_layer, len(attentions), time.time()-start)
-            start = time.time()
-            grad_matrix, attn_matrix = self.get_grad_attn_matrix(inputs_embeds, attentions, mean_attns, target_layer)
-            print(6,target_layer, len(attentions), time.time()-start)
-            start = time.time()
-            grad_matrices_all.append(grad_matrix.tolist())
+            #grad_matrix, attn_matrix = self.get_grad_attn_matrix(inputs_embeds, attentions, mean_attns, target_layer)
+            
+            attn_matrix = mean_attns[target_layer]
+            seq_len = attn_matrix.shape[0]
+            attn_matrix = attn_matrix[:seq_len, :seq_len]
             attn_matrices_all.append(attn_matrix.tolist())
-            print(7,target_layer, len(attentions), time.time()-start)
+
+
+
             start = time.time()
+            def scalar_outputs(inputs_embeds):
+
+                outputs = self.model.roberta(
+                    inputs_embeds=inputs_embeds,
+                    attention_mask=inputs["attention_mask"],
+                    output_attentions=True
+                )
+                attentions = outputs.attentions  
+                return attentions[target_layer].mean(dim=0).mean(dim=0).sum(dim=0)
+        
+            jac = torch.autograd.functional.jacobian(scalar_outputs, inputs_embeds).norm(dim=-1).squeeze(1)
+            
+            grad_matrices_all.append(jac.tolist())
+            print(1,time.time()-start)
+
+            start = time.time()
+            grad_norms_list = []
+
+            for k in range(seq_len):
+                scalar = attentions[target_layer].mean(dim=0).mean(dim=0)
+                scalar = scalar[:, k].sum()
+                
+                grad = torch.autograd.grad(scalar, inputs_embeds, retain_graph=True)[0].squeeze(0)
+                
+                grad_norms = grad.norm(dim=1)
+                grad_norms_list.append(grad_norms.unsqueeze(1))
+            print(2,time.time()-start)
 
-        print(8,time.time()-startloop)
         return grad_matrices_all, attn_matrices_all
     
     def get_grad_attn_matrix(self,inputs_embeds, attentions, mean_attns, target_layer):
         
-        start = time.time()
         attn_matrix = mean_attns[target_layer]
         seq_len = attn_matrix.shape[0]
-        
-        attn_matrix = torch.round(attn_matrix.float() * 100) / 100
-        attn_matrix = attn_matrix.to(torch.float16)
-        
         attn_layer = attentions[target_layer].squeeze(0).mean(dim=0)  # [seq, seq]
-        print(9,time.time()-start)
-        start = time.time()
-        #print('Computing grad norms')
+        """
+        print('Computing grad norms')
         grad_norms_list = []
+
         for k in range(seq_len):
             scalar = attn_layer[:, k].sum()
-            grad = torch.autograd.grad(scalar, inputs_embeds, retain_graph=True)[0].squeeze(0)
-            grad_norms = grad.norm(dim=1)
-
             
-            grad_norms = torch.round(grad_norms.unsqueeze(1).float() * 100) / 100
-            grad_norms = grad_norms.to(torch.float16)
+            grad = torch.autograd.grad(scalar, inputs_embeds, retain_graph=True)[0].squeeze(0)
             
+            grad_norms = grad.norm(dim=1)
+            grad_norms_list.append(grad_norms.unsqueeze(1))
         
-            grad_norms_list.append(grad_norms)
-        
-        print(10,time.time()-start)
-        start = time.time()
         grad_matrix = torch.cat(grad_norms_list, dim=1)
-        grad_matrix = grad_matrix[:seq_len, :seq_len]
-        attn_matrix = attn_matrix[:seq_len, :seq_len]
-        
-        print(11,time.time()-start)
-        start = time.time()
         
-        attn_matrix = torch.round(attn_matrix.float() * 100) / 100
-        attn_matrix = attn_matrix.to(torch.float16)
 
-        grad_matrix = torch.round(grad_matrix.float() * 100) / 100
-        grad_matrix = grad_matrix.to(torch.float16)
-        print(12,time.time()-start)
 
-        
 
+        grad_matrix = grad_matrix[:seq_len, :seq_len]
+        """
+        attn_matrix = attn_matrix[:seq_len, :seq_len]
+        grad_matrix = attn_matrix
         
         return grad_matrix, attn_matrix