Update ROBERTAmodel.py

ROBERTAmodel.py

@@ -183,6 +183,42 @@ class RoBERTaVisualizer(TransformerVisualizer):
         print('Average attentions per layer')
         mean_attns = [a.squeeze(0).mean(dim=0).detach().cpu() for a in attentions]
 
+
+        
+
+        def scalar_outputs(inputs_embeds):
+
+            outputs = self.model.roberta(
+                inputs_embeds=inputs_embeds,
+                attention_mask=inputs["attention_mask"],
+                output_attentions=True
+            )
+            attentions_condensed = [a.mean(dim=0).mean(dim=0).sum(dim=0) for a in attentions]
+            print([a.shape for a in attentions_condensed])
+            attentions_condensed= torch.vstack(attentions_condensed)
+            print(attentions_condensed.shape) 
+            return attentions_condensed
+    
+        start = time.time()
+        jac = torch.autograd.functional.jacobian(scalar_outputs, inputs_embeds).norm(dim=-1).squeeze(dim=2)
+        print(jac.shape)
+        grad_matrices_all = [jac[i] for i in range(jac.size(0))]
+
+        print(31,time.time()-start)
+        attn_matrices_all = []
+        for target_layer in range(len(attentions)):
+            #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(3,time.time()-start)
+ 
+
+
+        """
+
         attn_matrices_all = []
         grad_matrices_all = []
         for target_layer in range(len(attentions)):
@@ -194,7 +230,7 @@ class RoBERTaVisualizer(TransformerVisualizer):
             attn_matrices_all.append(attn_matrix.tolist())
 
 
-            """
+            
             start = time.time()
             def scalar_outputs(inputs_embeds):
 
@@ -210,7 +246,8 @@ 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)
@@ -221,6 +258,7 @@ class RoBERTaVisualizer(TransformerVisualizer):
                 grad_norms = grad.norm(dim=1)
                 grad_norms_list.append(grad_norms.unsqueeze(1))
             print(2,time.time()-start)
+        """
 
         return grad_matrices_all, attn_matrices_all