Update ROBERTAmodel.py

ROBERTAmodel.py

@@ -186,16 +186,23 @@ class RoBERTaVisualizer(TransformerVisualizer):
         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())
             attn_matrices_all.append(attn_matrix.tolist())
+            print(7,target_layer, len(attentions), time.time()-start)
+            start = time.time()
 
-        print(5,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):
@@ -209,12 +216,10 @@ 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)
@@ -224,7 +229,6 @@ class RoBERTaVisualizer(TransformerVisualizer):
             grad_norms = grad_norms.to(torch.float16)
             
             start = time.time()
-            print(8,k,time.time()-start)
         
             grad_norms_list.append(grad_norms)
         
@@ -234,7 +238,6 @@ class RoBERTaVisualizer(TransformerVisualizer):
         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)
@@ -242,7 +245,6 @@ class RoBERTaVisualizer(TransformerVisualizer):
         grad_matrix = torch.round(grad_matrix.float() * 100) / 100
         grad_matrix = grad_matrix.to(torch.float16)
         start = time.time()
-        print(11,time.time()-start)