make flatten coordinates more generic; add forward_kwargs option; image tensors (#64)

Author: mschrimpf

model_tools/activations/core.py

@@ -97,7 +97,6 @@ def _expand_paths(self, activations, original_paths):
         index = [argsort_indices[i] for i in sorted_index]
         return activations[{'stimulus_path': index}]
 
-
     def register_batch_activations_hook(self, hook):
         r"""
         The hook will be called every time a batch of activations is retrieved.
@@ -195,26 +194,32 @@ def _package(self, layer_activations, stimuli_paths):
     def _package_layer(self, layer_activations, layer, stimuli_paths):
         assert layer_activations.shape[0] == len(stimuli_paths)
         activations, flatten_indices = flatten(layer_activations, return_index=True)  # collapse for single neuroid dim
-        assert flatten_indices.shape[1] in [1, 2, 3]
-        # see comment in _package for an explanation why we cannot simply have 'channel' for the FC layer
-        if flatten_indices.shape[1] == 1:    # FC
+        flatten_coord_names = None
+        if flatten_indices.shape[1] == 1:  # fully connected, e.g. classifier
+            # see comment in _package for an explanation why we cannot simply have 'channel' for the FC layer
             flatten_coord_names = ['channel', 'channel_x', 'channel_y']
-        elif flatten_indices.shape[1] == 2:  # Transformer
+        elif flatten_indices.shape[1] == 2:  # Transformer, e.g. ViT
             flatten_coord_names = ['channel', 'embedding']
-        elif flatten_indices.shape[1] == 3:  # 2DConv
+        elif flatten_indices.shape[1] == 3:  # 2DConv, e.g. resnet
             flatten_coord_names = ['channel', 'channel_x', 'channel_y']
-        flatten_coords = {flatten_coord_names[i]: [sample_index[i] if i < flatten_indices.shape[1] else np.nan for sample_index in flatten_indices]
-                          for i in range(len(flatten_coord_names))}
-        layer_assembly = NeuroidAssembly(
-            activations,
-            coords={**{'stimulus_path': stimuli_paths,
-                       'neuroid_num': ('neuroid', list(range(activations.shape[1]))),
-                       'model': ('neuroid', [self.identifier] * activations.shape[1]),
-                       'layer': ('neuroid', [layer] * activations.shape[1]),
-                       },
-                    **{coord: ('neuroid', values) for coord, values in flatten_coords.items()}},
-            dims=['stimulus_path', 'neuroid']
-        )
+        elif flatten_indices.shape[1] == 4:  # temporal sliding window, e.g. omnivron
+            flatten_coord_names = ['channel_temporal', 'channel_x', 'channel_y', 'channel']
+        else:
+            # we still package the activations, but are unable to provide channel information
+            self._logger.debug(f"Unknown layer activations shape {layer_activations.shape}, not inferring channels")
+
+        # build assembly
+        coords = {'stimulus_path': stimuli_paths,
+                  'neuroid_num': ('neuroid', list(range(activations.shape[1]))),
+                  'model': ('neuroid', [self.identifier] * activations.shape[1]),
+                  'layer': ('neuroid', [layer] * activations.shape[1]),
+                  }
+        if flatten_coord_names:
+            flatten_coords = {flatten_coord_names[i]: [sample_index[i] if i < flatten_indices.shape[1] else np.nan
+                                                       for sample_index in flatten_indices]
+                              for i in range(len(flatten_coord_names))}
+            coords = {**coords, **{coord: ('neuroid', values) for coord, values in flatten_coords.items()}}
+        layer_assembly = NeuroidAssembly(activations, coords=coords, dims=['stimulus_path', 'neuroid'])
         neuroid_id = [".".join([f"{value}" for value in values]) for values in zip(*[
             layer_assembly[coord].values for coord in ['model', 'layer', 'neuroid_num']])]
         layer_assembly['neuroid_id'] = 'neuroid', neuroid_id

model_tools/activations/pytorch.py

@@ -1,6 +1,6 @@
-import logging
 from collections import OrderedDict
 
+import logging
 import numpy as np
 from PIL import Image
 
@@ -11,7 +11,7 @@
 
 
 class PytorchWrapper:
-    def __init__(self, model, preprocessing, identifier=None, *args, **kwargs):
+    def __init__(self, model, preprocessing, identifier=None, forward_kwargs=None, *args, **kwargs):
         import torch
         logger = logging.getLogger(fullname(self))
         self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -22,6 +22,7 @@ def __init__(self, model, preprocessing, identifier=None, *args, **kwargs):
         self._extractor = self._build_extractor(
             identifier=identifier, preprocessing=preprocessing, get_activations=self.get_activations, *args, **kwargs)
         self._extractor.insert_attrs(self)
+        self._forward_kwargs = forward_kwargs or {}
 
     def _build_extractor(self, identifier, preprocessing, get_activations, *args, **kwargs):
         return ActivationsExtractorHelper(
@@ -42,7 +43,7 @@ def __call__(self, *args, **kwargs):  # cannot assign __call__ as attribute due
     def get_activations(self, images, layer_names):
         import torch
         from torch.autograd import Variable
-        images = [torch.from_numpy(image) for image in images]
+        images = [torch.from_numpy(image) if not isinstance(image, torch.Tensor) else image for image in images]
         images = Variable(torch.stack(images))
         images = images.to(self._device)
         self._model.eval()
@@ -55,7 +56,8 @@ def get_activations(self, images, layer_names):
             hook = self.register_hook(layer, layer_name, target_dict=layer_results)
             hooks.append(hook)
 
-        self._model(images)
+        with torch.no_grad():
+            self._model(images, **self._forward_kwargs)
         for hook in hooks:
             hook.remove()
         return layer_results
@@ -115,7 +117,7 @@ def load_images(image_filepaths):
 
 def load_image(image_filepath):
     with Image.open(image_filepath) as pil_image:
-        if 'L' not in pil_image.mode.upper() and 'A' not in pil_image.mode.upper()\
+        if 'L' not in pil_image.mode.upper() and 'A' not in pil_image.mode.upper() \
                 and 'P' not in pil_image.mode.upper():  # not binary and not alpha and not palletized
             # work around to https://github.com/python-pillow/Pillow/issues/1144,
             # see https://stackoverflow.com/a/30376272/2225200