Update ROBERTAmodel.py

ROBERTAmodel.py

@@ -200,6 +200,7 @@ class RoBERTaVisualizer(TransformerVisualizer):
     
     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]
         
@@ -208,9 +209,12 @@ class RoBERTaVisualizer(TransformerVisualizer):
         
         attn_layer = attentions[target_layer].squeeze(0).mean(dim=0)  # [seq, seq]
 
+        print(6,time.time()-start)
         print('Computing grad norms')
         grad_norms_list = []
         for k in range(seq_len):
+            start = time.time()
+            print(7,k,time.time()-start)
             scalar = attn_layer[:, k].sum()
             grad = torch.autograd.grad(scalar, inputs_embeds, retain_graph=True)[0].squeeze(0)
             grad_norms = grad.norm(dim=1)
@@ -218,18 +222,27 @@ class RoBERTaVisualizer(TransformerVisualizer):
             
             grad_norms = torch.round(grad_norms.unsqueeze(1).float() * 100) / 100
             grad_norms = grad_norms.to(torch.float16)
+            
+            start = time.time()
+            print(8,k,time.time()-start)
         
             grad_norms_list.append(grad_norms)
-
+        
+        start = time.time()
+        print(9,time.time()-start)
         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]
+        start = time.time()
+        print(10,time.time()-start)
         
         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)
+        start = time.time()
+        print(11,time.time()-start)