Update ROBERTAmodel.py

ROBERTAmodel.py

@@ -194,7 +194,7 @@ class RoBERTaVisualizer(TransformerVisualizer):
             attn_matrices_all.append(attn_matrix.tolist())
 
 
-
+            """
             start = time.time()
             def scalar_outputs(inputs_embeds):
 
@@ -210,14 +210,12 @@ class RoBERTaVisualizer(TransformerVisualizer):
             
             grad_matrices_all.append(jac.tolist())
             print(1,time.time()-start)
-
+            """
             start = time.time()
             grad_norms_list = []
-
+            scalar_layer = attentions[target_layer].mean(dim=0).mean(dim=0)
             for k in range(seq_len):
-                scalar = attentions[target_layer].mean(dim=0).mean(dim=0)
-                scalar = scalar[:, k].sum()
-                
+                scalar = scalar_layer[:, k].sum()
                 grad = torch.autograd.grad(scalar, inputs_embeds, retain_graph=True)[0].squeeze(0)
                 
                 grad_norms = grad.norm(dim=1)
@@ -225,32 +223,4 @@ class RoBERTaVisualizer(TransformerVisualizer):
             print(2,time.time()-start)
 
         return grad_matrices_all, attn_matrices_all
-    
-    def get_grad_attn_matrix(self,inputs_embeds, attentions, mean_attns, target_layer):
-        
-        attn_matrix = mean_attns[target_layer]
-        seq_len = attn_matrix.shape[0]
-        attn_layer = attentions[target_layer].squeeze(0).mean(dim=0)  # [seq, seq]
-        """
-        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_list.append(grad_norms.unsqueeze(1))
-        
-        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]
-        grad_matrix = attn_matrix
-        
-        return grad_matrix, attn_matrix
+