brainscore-vision 2.1__py3-none-any.whl

brainscore_vision/model_helpers/activations/pca.py

@@ -0,0 +1,117 @@
+import h5py
+import logging
+import numpy as np
+import os
+from PIL import Image
+from result_caching import store_dict
+from sklearn.decomposition import PCA
+from tqdm import tqdm
+
+from brainscore_vision.model_helpers.activations.core import flatten, change_dict
+from brainscore_vision.model_helpers.utils import fullname, s3
+
+
+class LayerPCA:
+    def __init__(self, activations_extractor, n_components):
+        self._logger = logging.getLogger(fullname(self))
+        self._extractor = activations_extractor
+        self._n_components = n_components
+        self._layer_pcas = {}
+
+    def __call__(self, batch_activations):
+        self._ensure_initialized(batch_activations.keys())
+
+        def apply_pca(layer, activations):
+            pca = self._layer_pcas[layer]
+            activations = flatten(activations)
+            if pca is None:
+                return activations
+            return pca.transform(activations)
+
+        return change_dict(batch_activations, apply_pca, keep_name=True,
+                           multithread=os.getenv('MT_MULTITHREAD', '1') == '1')
+
+    def _ensure_initialized(self, layers):
+        missing_layers = [layer for layer in layers if layer not in self._layer_pcas]
+        if len(missing_layers) == 0:
+            return
+        layer_pcas = self._pcas(identifier=self._extractor.identifier, layers=missing_layers,
+                                n_components=self._n_components)
+        self._layer_pcas = {**self._layer_pcas, **layer_pcas}
+
+    @store_dict(dict_key='layers', identifier_ignore=['layers'])
+    def _pcas(self, identifier, layers, n_components):
+        self._logger.debug('Retrieving ImageNet activations')
+        imagenet_paths = _get_imagenet_val(num_images=n_components)
+        self.handle.disable()
+        imagenet_activations = self._extractor(imagenet_paths, layers=layers)
+        imagenet_activations = {layer: imagenet_activations.sel(layer=layer).values
+                                for layer in np.unique(imagenet_activations['layer'])}
+        assert len(set(activations.shape[0] for activations in imagenet_activations.values())) == 1, "stimuli differ"
+        self.handle.enable()
+
+        self._logger.debug('Computing ImageNet principal components')
+        progress = tqdm(total=len(imagenet_activations), desc="layer principal components")
+
+        def init_and_progress(layer, activations):
+            activations = flatten(activations)
+            if activations.shape[1] <= n_components:
+                self._logger.debug(f"Not computing principal components for {layer} "
+                                   f"activations {activations.shape} as shape is small enough already")
+                pca = None
+            else:
+                pca = PCA(n_components=n_components, random_state=0)
+                pca.fit(activations)
+            progress.update(1)
+            return pca
+
+        layer_pcas = change_dict(imagenet_activations, init_and_progress, keep_name=True,
+                                 multithread=os.getenv('MT_MULTITHREAD', '1') == '1')
+        progress.close()
+        return layer_pcas
+
+    @classmethod
+    def hook(cls, activations_extractor, n_components):
+        hook = LayerPCA(activations_extractor=activations_extractor, n_components=n_components)
+        assert not cls.is_hooked(activations_extractor), "PCA already hooked"
+        handle = activations_extractor.register_batch_activations_hook(hook)
+        hook.handle = handle
+        return handle
+
+    @classmethod
+    def is_hooked(cls, activations_extractor):
+        return any(isinstance(hook, cls) for hook in
+                   activations_extractor._extractor._batch_activations_hooks.values())
+
+
+def _get_imagenet_val(num_images):
+    _logger = logging.getLogger(fullname(_get_imagenet_val))
+    num_classes = 1000
+    num_images_per_class = (num_images - 1) // num_classes
+    base_indices = np.arange(num_images_per_class).astype(int)
+    indices = []
+    for i in range(num_classes):
+        indices.extend(50 * i + base_indices)
+    for i in range((num_images - 1) % num_classes + 1):
+        indices.extend(50 * i + np.array([num_images_per_class]).astype(int))
+
+    framework_home = os.path.expanduser(os.getenv('MT_HOME', '~/.model-tools'))
+    imagenet_filepath = os.getenv('MT_IMAGENET_PATH', os.path.join(framework_home, 'imagenet2012.hdf5'))
+    imagenet_dir = f"{imagenet_filepath}-files"
+    os.makedirs(imagenet_dir, exist_ok=True)
+
+    if not os.path.isfile(imagenet_filepath):
+        os.makedirs(os.path.dirname(imagenet_filepath), exist_ok=True)
+        _logger.debug(f"Downloading ImageNet validation to {imagenet_filepath}")
+        s3.download_file("imagenet2012-val.hdf5", imagenet_filepath)
+
+    filepaths = []
+    with h5py.File(imagenet_filepath, 'r') as f:
+        for index in indices:
+            imagepath = os.path.join(imagenet_dir, f"{index}.png")
+            if not os.path.isfile(imagepath):
+                image = np.array(f['val/images'][index])
+                Image.fromarray(image).save(imagepath)
+            filepaths.append(imagepath)
+
+    return filepaths

brainscore_vision/model_helpers/activations/pytorch.py

@@ -0,0 +1,152 @@
+from collections import OrderedDict
+
+import logging
+import numpy as np
+from PIL import Image
+
+from brainscore_vision.model_helpers.activations.core import ActivationsExtractorHelper
+from brainscore_vision.model_helpers.utils import fullname
+
+SUBMODULE_SEPARATOR = '.'
+
+
+class PytorchWrapper:
+    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")
+        logger.debug(f"Using device {self._device}")
+        self._model = model
+        self._model = self._model.to(self._device)
+        identifier = identifier or model.__class__.__name__
+        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(
+            identifier=identifier, get_activations=get_activations, preprocessing=preprocessing,
+            *args, **kwargs)
+
+    @property
+    def identifier(self):
+        return self._extractor.identifier
+
+    @identifier.setter
+    def identifier(self, value):
+        self._extractor.identifier = value
+
+    def __call__(self, *args, **kwargs):  # cannot assign __call__ as attribute due to Python convention
+        return self._extractor(*args, **kwargs)
+
+    def get_activations(self, images, layer_names):
+        import torch
+        from torch.autograd import Variable
+        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()
+
+        layer_results = OrderedDict()
+        hooks = []
+
+        for layer_name in layer_names:
+            layer = self.get_layer(layer_name)
+            hook = self.register_hook(layer, layer_name, target_dict=layer_results)
+            hooks.append(hook)
+
+        with torch.no_grad():
+            self._model(images, **self._forward_kwargs)
+        for hook in hooks:
+            hook.remove()
+        return layer_results
+
+    def get_layer(self, layer_name):
+        if layer_name == 'logits':
+            return self._output_layer()
+        module = self._model
+        for part in layer_name.split(SUBMODULE_SEPARATOR):
+            module = module._modules.get(part)
+            assert module is not None, f"No submodule found for layer {layer_name}, at part {part}"
+        return module
+
+    def _output_layer(self):
+        module = self._model
+        while module._modules:
+            module = module._modules[next(reversed(module._modules))]
+        return module
+
+    @classmethod
+    def _tensor_to_numpy(cls, output):
+        return output.cpu().data.numpy()
+
+    def register_hook(self, layer, layer_name, target_dict):
+        def hook_function(_layer, _input, output, name=layer_name):
+            target_dict[name] = PytorchWrapper._tensor_to_numpy(output)
+
+        hook = layer.register_forward_hook(hook_function)
+        return hook
+
+    def __repr__(self):
+        return repr(self._model)
+
+    def layers(self):
+        for name, module in self._model.named_modules():
+            if len(list(module.children())) > 0:  # this module only holds other modules
+                continue
+            yield name, module
+
+    def graph(self):
+        import networkx as nx
+        g = nx.DiGraph()
+        for layer_name, layer in self.layers():
+            g.add_node(layer_name, object=layer, type=type(layer))
+        return g
+
+
+def load_preprocess_images(image_filepaths, image_size, **kwargs):
+    images = load_images(image_filepaths)
+    images = preprocess_images(images, image_size=image_size, **kwargs)
+    return images
+
+
+def load_images(image_filepaths):
+    return [load_image(image_filepath) for image_filepath in 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() \
+                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
+            return pil_image.copy()
+        else:  # make sure potential binary images are in RGB
+            rgb_image = Image.new("RGB", pil_image.size)
+            rgb_image.paste(pil_image)
+            return rgb_image
+
+
+def preprocess_images(images, image_size, **kwargs):
+    preprocess = torchvision_preprocess_input(image_size, **kwargs)
+    images = [preprocess(image) for image in images]
+    images = np.concatenate(images)
+    return images
+
+
+def torchvision_preprocess_input(image_size, **kwargs):
+    from torchvision import transforms
+    return transforms.Compose([
+        transforms.Resize((image_size, image_size)),
+        torchvision_preprocess(**kwargs),
+    ])
+
+
+def torchvision_preprocess(normalize_mean=(0.485, 0.456, 0.406), normalize_std=(0.229, 0.224, 0.225)):
+    from torchvision import transforms
+    return transforms.Compose([
+        transforms.ToTensor(),
+        transforms.Normalize(mean=normalize_mean, std=normalize_std),
+        lambda img: img.unsqueeze(0)
+    ])

brainscore_vision/model_helpers/activations/temporal/core/__init__.py

@@ -0,0 +1,3 @@
+from .extractor import ActivationsExtractor
+from .executor import BatchExecutor
+from .inferencer import *

brainscore_vision/model_helpers/activations/temporal/core/executor.py

@@ -0,0 +1,219 @@
+import os
+import logging
+
+import numpy as np
+from tqdm.auto import tqdm
+from collections import OrderedDict
+from typing import Any, Callable, Dict, Hashable, List
+from ..inputs import Stimulus
+
+from brainscore_vision.model_helpers.utils import fullname
+from joblib import Parallel, delayed
+
+
+# a utility to apply a list of functions to inputs sequentially but only iterate over the first input
+def _pipeline(*funcs):
+    def _func(x, *others):
+        for f in funcs:
+            x = f(x, *others)
+        return x
+    return _func
+
+    
+# a mapper that execute a function in parallel with joblib
+class JoblibMapper:
+    def __init__(self, num_threads: int):
+        self._num_threads = num_threads
+        self._pool = Parallel(n_jobs=num_threads, verbose=False, backend="loky")
+        self._failed_to_pickle_func = False
+
+    def map(self, func, *data):
+        from joblib.externals.loky.process_executor import TerminatedWorkerError, BrokenProcessPool
+        if not self._failed_to_pickle_func:
+            try:
+                return self._pool(delayed(func)(*x) for x in zip(*data))
+            except (TerminatedWorkerError, BrokenProcessPool):
+                self._failed_to_pickle_func = True
+        return [func(*x) for x in zip(*data)]
+
+
+class BatchExecutor:
+    """Executor for batch processing of stimuli.
+    
+    Parameters
+    ----------
+        get_activations : function 
+            function that takes a list of processed stimuli and a list of layers, and returns a dictionary of activations.
+        preprocessing : function
+            function that takes a stimulus and returns a processed stimulus.
+        batch_size: int
+            number of stimuli to process in each batch.
+        batch_padding: bool
+            whether to pad the each batch with the last stimulus to make it the same size as the specified batch size. 
+            Otherwise, some batches will have size < batch_size because of the lacking of samples in that group.
+        batch_grouper: function
+            function that takes a stimulus and return the property based on which the stimuli can be grouped.
+        max_workers: int
+            number of workers for parallel processing. If None, the number of workers will be the number of cpus.
+
+    APIs
+    ----
+    add_stimuli(stimuli)
+        add a list of stimulus to the executor.
+    clear_stimuli()
+        clear the kept stimuli list.
+    execute(layers)
+        execute the batch processing of the kept stimuli and return the activations of the specified layers.
+        For each layer, the activation of different stimuli will be stacked.
+        If they have different shapes, the smaller ones will be padded with NaNs.
+
+    Examples
+    --------
+    >>> executor = BatchExecutor(get_activations, preprocessing, batch_size=64, batch_grouper=lambda s: s.size)
+    >>> executor.register_before_hook(lambda v: v.set_size(180, 120))
+    >>> executor.register_after_hook(lambda val, l, s: val if l == "layer1" else val.mean(0))
+    >>> executor.add_stimuli(stimulus_list_1)
+    >>> executor.add_stimuli(stimulus_list_2)
+    >>> layer_activations = executor.execute(layers)  # here the actual batching and processing happen
+    """
+
+
+    def __init__(self, 
+                get_activations : Callable[[List[Any]], Dict[str, np.array]], 
+                preprocessing : Callable[[List[Stimulus]], Any], 
+                batch_size : int, 
+                batch_padding : bool = False, 
+                batch_grouper : Callable[[Stimulus], Hashable] = None, 
+                max_workers : int = None
+            ):
+        self.stimuli = []
+        self.get_activations = get_activations
+        self.batch_size = batch_size
+        self.batch_padding = batch_padding
+        self.batch_grouper = batch_grouper
+        self.preprocess = preprocessing
+        self.max_workers = max_workers
+
+        self._logger = logging.getLogger(fullname(self))
+
+        # Pool for I/O intensive ops
+        # torch suggest using the number of cpus as the number of threads, but os.cpu_count() returns the number of threads
+        num_threads = max(min(int(os.cpu_count() / 1.5), self.batch_size), 1)  
+        if self.max_workers is not None:
+            num_threads = min(self.max_workers, num_threads)
+        self._logger.info(f"Using {num_threads} threads for parallel processing.")
+        self._mapper = JoblibMapper(num_threads)
+
+        # processing hooks
+        self.before_hooks = []
+        self.after_hooks = []
+
+    def _get_batches(
+            self, 
+            data, 
+            batch_size : int, 
+            grouper : Callable[[Stimulus], Hashable] = None, 
+            padding : bool = False
+        ):
+        """Group the data into batches based on the grouper.
+        
+        Parameters
+        ----------
+
+        data : array-like
+            list of data to be grouped.
+        batch_size, grouper, padding : int, function, bool, directly set by the class
+
+        Returns
+        -------
+        indices : list
+            indices of the source data after sorting.
+        masks : list of list
+            masks for each batch to indicate whether the datum is padding sample.
+        all_batches : list of list
+            list of batches.
+        """
+
+        N = len(data)
+
+        if grouper is None:
+            sorted_data = np.array(data, dtype='object')
+            inverse_indices = np.arange(N)
+            sorted_properties = [0] * N
+        else:
+            properties = np.array([hash(grouper(d)) for d in data])
+            sorted_indices = np.argsort(properties)
+            inverse_indices = np.argsort(sorted_indices)  # inverse transform
+            sorted_properties = properties[sorted_indices]
+            sorted_data = np.array(data, dtype='object')[sorted_indices]
+        inverse_indices = list(inverse_indices)
+
+        index = 0
+        all_batches = []
+        all_indices = []
+        indices = []
+        masks = []
+        while index < N:
+            property_anchor = sorted_properties[index]
+            batch = []
+            while index < N and len(batch) < batch_size and sorted_properties[index] == property_anchor:
+                batch.append(sorted_data[index])
+                index += 1
+            
+            batch_indices = inverse_indices[index-len(batch):index]
+
+            if padding:
+                num_padding = batch_size - len(batch)
+                if num_padding:
+                    batch_padding = [batch[-1]] * num_padding
+                    batch += batch_padding
+            else:
+                num_padding = 0
+            masks.append([True] * (len(batch)-num_padding) + [False] * num_padding) 
+
+            all_batches.append(batch)
+            all_indices.append(batch_indices)
+            indices.extend(batch_indices)
+        return indices, masks, all_batches
+
+    def register_before_hook(self, hook: Callable[[Stimulus], Stimulus], index=None):
+        # hook: Stimulus -> Stimulus
+        if index is None: index = len(self.before_hooks)
+        self.before_hooks.insert(index, hook)
+
+    def register_after_hook(self, hook: Callable[[Any, str, Stimulus], Any], index=None):
+        # hook: value, layer_name, Stimulus -> value
+        if index is None: index = len(self.after_hooks)
+        self.after_hooks.insert(index, hook)
+    
+    def add_stimuli(self, stimuli):
+        self.stimuli.extend(stimuli)
+
+    def clear_stimuli(self):
+        self.stimuli = []
+
+    def execute(self, layers):
+        indices, masks, batches = self._get_batches(self.stimuli, self.batch_size, 
+                                                    grouper=self.batch_grouper,
+                                                    padding=self.batch_padding)
+        
+        before_pipe = _pipeline(*self.before_hooks)
+        after_pipe = _pipeline(*self.after_hooks)
+        
+        layer_activations = OrderedDict()
+        for mask, batch in tqdm(zip(masks, batches), desc="activations", total=len(batches)):
+            batch = [before_pipe(stimulus) for stimulus in batch]
+            model_inputs = self._mapper.map(self.preprocess, batch)
+            batch_activations = self.get_activations(model_inputs, layers)
+            assert isinstance(batch_activations, OrderedDict)
+            for layer, activations in batch_activations.items():
+                results = [after_pipe(arr, layer, stimulus) 
+                               for not_pad, arr, stimulus in zip(mask, activations, batch) 
+                               if not_pad]
+                layer_activations.setdefault(layer, []).extend(results)
+
+        for layer, activations in layer_activations.items():
+            layer_activations[layer] = [activations[i] for i in indices]
+
+        self.clear_stimuli()
+        return layer_activations