get cls block

.gitignore

@@ -0,0 +1,2 @@
+__pycache__
+dino.egg-info

vision_transformer.py

@@ -21,7 +21,7 @@
 import torch
 import torch.nn as nn
 
-from utils import trunc_normal_
+from .utils import trunc_normal_
 
 
 def drop_path(x, drop_prob: float = 0., training: bool = False):
@@ -213,6 +213,13 @@ def forward(self, x):
         x = self.norm(x)
         return x[:, 0]
 
+    def get_last_block(self, x):
+        x = self.prepare_tokens(x)
+        for blk in self.blocks:
+            x = blk(x)
+        x = self.norm(x)
+        return x[0]
+
     def get_last_selfattention(self, x):
         x = self.prepare_tokens(x)
         for i, blk in enumerate(self.blocks):

visualize_attention.py

@@ -32,7 +32,8 @@
 from PIL import Image
 
 import utils
-import vision_transformer as vits
+
+from . import vision_transformer as vits
 
 
 def apply_mask(image, mask, color, alpha=0.5):