added only extraction of class token as an option

ViCCo-Group · Aug 3, 2023 · 3079f3c · 3079f3c
1 parent 7245361
commit 3079f3c
Showing 1 changed file with 10 additions and 0 deletions.
diff --git a/thingsvision/core/extraction/torch.py b/thingsvision/core/extraction/torch.py
@@ -31,6 +31,11 @@ def __init__(
         self.activations = {}
         self.hook_handle = None
 
+        if 'extract_cls_token' in self.model_parameters:
+            self.extract_cls_token = self.model_parameters['extract_cls_token']
+        else:
+            self.extract_cls_token = False
+
         if not self.model:
             self.load_model()
         self.prepare_inference()
@@ -68,6 +73,11 @@ def hook(model, input, output) -> None:
                 act = output[0]
             else:
                 act = output
+
+            if self.extract_cls_token:
+                # we extract only the representations of the first token (cls token)
+                act = act[:, 0]
+
             try:
                 self.activations[name] = act.clone().detach()
             except AttributeError: