visualize_attention_in_ViT

Running

maomao88 commited on Jun 23, 2025

Commit

574e06d

1 Parent(s): c19dba1

more updates

Files changed (4) hide show

__pycache__/main.cpython-313.pyc CHANGED Viewed

Binary files a/__pycache__/main.cpython-313.pyc and b/__pycache__/main.cpython-313.pyc differ

main.py CHANGED Viewed

@@ -79,7 +79,7 @@ def compute_cls_rollout_and_raw_others(attentions, start_layer=0):
     cls_rollout = rollout[0]  # shape: [num_tokens]
     # Get raw attention from last layer (no rollout) for other tokens
-    last_attention = attentions[-1][0].mean(dim=0)  # [tokens, tokens]
     # Replace CLS row in last_attention with rollout
     last_attention[0] = cls_rollout
@@ -99,7 +99,7 @@ def attention_rollout(attentions, start_layer=0):
     for attention in attentions[start_layer:]:
         # Average over heads
-        attn_heads_fused = attention[-1].mean(dim=0)  # [tokens, tokens]
         # Add identity matrix to account for residual connections
         attn_heads_fused = attn_heads_fused + torch.eye(attn_heads_fused.size(0))

     cls_rollout = rollout[0]  # shape: [num_tokens]
     # Get raw attention from last layer (no rollout) for other tokens
+    last_attention = attentions[0][0].mean(dim=0)  # [tokens, tokens]
     # Replace CLS row in last_attention with rollout
     last_attention[0] = cls_rollout
     for attention in attentions[start_layer:]:
         # Average over heads
+        attn_heads_fused = attention[0].mean(dim=0)  # [tokens, tokens]
         # Add identity matrix to account for residual connections
         attn_heads_fused = attn_heads_fused + torch.eye(attn_heads_fused.size(0))

static/index.html CHANGED Viewed

@@ -8,28 +8,28 @@
 <body>
   <div class="container">
     <h2>ViT Image Classification with Attention Visualization</h2>
-      <div class="desc">
-          <p>Upload an image to classify its content and visualize the attention map on the right.</p>
-          <p>By default, the attention from the <b>[CLS]</b> token, representing the model’s overall summary, is highlighted, showing which image regions contributed most to the classification.</p>
-          <p>Click on any region of the image on the right to explore how the model attends to that specific area.</p>
-      </div>
     <div class="content-container">
         <div class="upload-container">
             <form id="upload-form">
-              <div id="drop-zone" class="drop-zone">
-                <span id="drop-zone-text">Drop image here or <button type="button" id="browse-btn">Browse</button></span>
-                <input type="file" id="image-input" accept="image/*" style="display: none;" required />
-              </div>
-              <button type="submit" id="submit-btn">Classify Image</button>
             </form>
             <div id="local-images" class="examples">
                 <span>Examples: </span>
             </div>
         </div>
        <div id="attn"></div>
     </div>
-    <div id="result"></div>
     <div id="preview"></div>
   </div>
   <script src="app.js"></script>

 <body>
   <div class="container">
     <h2>ViT Image Classification with Attention Visualization</h2>
+    <div class="desc">
+        <p>Upload an image to classify its content and visualize the attention map on the right.</p>
+        <p>By default, the attention from the <b>[CLS]</b> token, representing the model’s overall summary, is highlighted, showing which image regions contributed most to the classification.</p>
+        <p>Click on any region of the image on the right to explore how the model attends to that specific area.</p>
+    </div>
     <div class="content-container">
         <div class="upload-container">
             <form id="upload-form">
+                <div id="drop-zone" class="drop-zone">
+                    <span id="drop-zone-text">Drop image here or <button type="button" id="browse-btn">Browse</button></span>
+                    <input type="file" id="image-input" accept="image/*" style="display: none;" required />
+                </div>
+                <button type="submit" id="submit-btn">Classify Image</button>
             </form>
             <div id="local-images" class="examples">
                 <span>Examples: </span>
             </div>
+            <div id="result"></div>
         </div>
        <div id="attn"></div>
     </div>
     <div id="preview"></div>
   </div>
   <script src="app.js"></script>

static/style.css CHANGED Viewed

@@ -155,7 +155,10 @@ button[type="submit"]:disabled {
 }
 #result {
-  margin-top: 16px;
   font-weight: bold;
   color: #333;
 }

 }
 #result {
+  margin-top: 2.5rem;
   font-weight: bold;
   color: #333;
+  padding: 1.3rem 0;
+  border: solid 3px #f97316;
+  border-radius: 6px;
 }