brainscore-vision 2.2.3__py3-none-any.whl → 2.2.5__py3-none-any.whl

brainscore_vision/models/ReAlnet/model.py

@@ -0,0 +1,237 @@
+import math
+from collections import OrderedDict
+import torch
+from torch import nn
+from torchvision import transforms
+import torch.utils.model_zoo
+import os
+from torch.utils.data import Dataset, DataLoader
+import pandas as pd
+import numpy as np
+import torch.nn.functional as F
+import h5py
+import random
+import functools
+import torchvision.models
+from brainscore_vision.model_helpers.s3 import load_weight_file
+from brainscore_vision.model_helpers.activations.pytorch import PytorchWrapper
+from brainscore_vision.model_helpers.activations.pytorch import load_preprocess_images
+import json
+
+LAYERS = ['V1', 'V2', 'V4', 'IT', 'decoder.avgpool']
+
+class Flatten(nn.Module):
+    """
+    Helper module for flattening input tensor to 1-D for the use in Linear modules
+    """
+    def forward(self, x):
+        return x.view(x.size(0), -1)
+
+class Identity(nn.Module):
+    """
+    Helper module that stores the current tensor. Useful for accessing by name
+    """
+    def forward(self, x):
+        return x
+
+class CORblock_S(nn.Module):
+    scale = 4  # scale of the bottleneck convolution channels
+
+    def __init__(self, in_channels, out_channels, times=1):
+        super().__init__()
+        self.times = times
+
+        self.conv_input = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False)
+        self.skip = nn.Conv2d(out_channels, out_channels,
+                              kernel_size=1, stride=2, bias=False)
+        self.norm_skip = nn.BatchNorm2d(out_channels)
+
+        self.conv1 = nn.Conv2d(out_channels, out_channels * self.scale,
+                               kernel_size=1, bias=False)
+        self.nonlin1 = nn.ReLU(inplace=True)
+
+        self.conv2 = nn.Conv2d(out_channels * self.scale, out_channels * self.scale,
+                               kernel_size=3, stride=2, padding=1, bias=False)
+        self.nonlin2 = nn.ReLU(inplace=True)
+
+        self.conv3 = nn.Conv2d(out_channels * self.scale, out_channels,
+                               kernel_size=1, bias=False)
+        self.nonlin3 = nn.ReLU(inplace=True)
+
+        self.output = Identity()  # for an easy access to this block's output
+
+        # need BatchNorm for each time step for training to work well
+        for t in range(self.times):
+            setattr(self, f'norm1_{t}', nn.BatchNorm2d(out_channels * self.scale))
+            setattr(self, f'norm2_{t}', nn.BatchNorm2d(out_channels * self.scale))
+            setattr(self, f'norm3_{t}', nn.BatchNorm2d(out_channels))
+
+    def forward(self, inp):
+        x = self.conv_input(inp)
+        for t in range(self.times):
+            if t == 0:
+                skip = self.norm_skip(self.skip(x))
+                self.conv2.stride = (2, 2)
+            else:
+                skip = x
+                self.conv2.stride = (1, 1)
+
+            x = self.conv1(x)
+            x = getattr(self, f'norm1_{t}')(x)
+            x = self.nonlin1(x)
+
+            x = self.conv2(x)
+            x = getattr(self, f'norm2_{t}')(x)
+            x = self.nonlin2(x)
+
+            x = self.conv3(x)
+            x = getattr(self, f'norm3_{t}')(x)
+
+            x += skip
+            x = self.nonlin3(x)
+            output = self.output(x)
+
+        return output
+
+def CORnet_S():
+    model = nn.Sequential(OrderedDict([
+        ('V1', nn.Sequential(OrderedDict([
+            ('conv1', nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)),
+            ('norm1', nn.BatchNorm2d(64)),
+            ('nonlin1', nn.ReLU(inplace=True)),
+            ('pool', nn.MaxPool2d(kernel_size=3, stride=2, padding=1)),
+            ('conv2', nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, bias=False)),
+            ('norm2', nn.BatchNorm2d(64)),
+            ('nonlin2', nn.ReLU(inplace=True)),
+            ('output', Identity())
+        ]))),
+        ('V2', CORblock_S(64, 128, times=2)),
+        ('V4', CORblock_S(128, 256, times=4)),
+        ('IT', CORblock_S(256, 512, times=2)),
+        ('decoder', nn.Sequential(OrderedDict([
+            ('avgpool', nn.AdaptiveAvgPool2d(1)),
+            ('flatten', Flatten()),
+            ('linear', nn.Linear(512, 1000)),
+            ('output', Identity())
+        ])))
+    ]))
+
+    # weight initialization
+    for m in model.modules():
+        if isinstance(m, nn.Conv2d):
+            n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
+            m.weight.data.normal_(0, math.sqrt(2. / n))
+        elif isinstance(m, nn.BatchNorm2d):
+            m.weight.data.fill_(1)
+            m.bias.data.zero_()
+
+    return model
+
+class Encoder(nn.Module):
+    def __init__(self, realnet, n_output):
+        super(Encoder, self).__init__()
+        
+        # CORnet
+        self.realnet = realnet
+        
+        # fully connected layers
+        self.fc_v1 = nn.Linear(200704, 128)
+        self.fc_v2 = nn.Linear(100352, 128)
+        self.fc_v4 = nn.Linear(50176, 128)
+        self.fc_it = nn.Linear(25088, 128)
+        self.fc = nn.Linear(512, n_output)
+        self.activation = nn.ReLU()
+        
+    def forward(self, imgs):
+        # forward pass through CORnet_S
+        outputs = self.realnet(imgs)
+        
+        N = len(imgs)
+        v1_outputs = self.realnet.V1(imgs)   # N * 64 * 56 * 56
+        v2_outputs = self.realnet.V2(v1_outputs)   # N * 128 * 28 * 28
+        v4_outputs = self.realnet.V4(v2_outputs)   # N * 256 * 14 * 14
+        it_outputs = self.realnet.IT(v4_outputs)   # N * 512 * 7 * 7
+
+        # flatten and pass through fully connected layers
+        v1_features = self.fc_v1(v1_outputs.view(N, -1))
+        v1_features = self.activation(v1_features)
+        
+        v2_features = self.fc_v2(v2_outputs.view(N, -1))
+        v2_features = self.activation(v2_features)
+        
+        v4_features = self.fc_v4(v4_outputs.view(N, -1))
+        v4_features = self.activation(v4_features)
+        
+        it_features = self.fc_it(it_outputs.view(N, -1))
+        it_features = self.activation(it_features)
+        
+        features = torch.cat((v1_features, v2_features, v4_features, it_features), dim=1)
+        features = self.fc(features)
+        
+        return outputs, features
+
+# Change here: use 'cpu'
+device = 'cpu'
+
+transform = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+])
+
+# Construct CORnet_S
+realnet = CORnet_S()
+# (Optional) remove DataParallel if not needed for CPU
+# realnet = torch.nn.DataParallel(realnet)
+
+def load_config(json_file):
+    # Get the directory containing this script (model.py)
+    base_dir = os.path.dirname(__file__)
+    
+    # Construct the path to the JSON file
+    json_path = os.path.join(base_dir, json_file)
+    
+    # Read the JSON
+    with open(json_path, "r", encoding="utf-8") as f:
+        data = json.load(f)
+    return data
+
+
+# Build encoder model
+encoder = Encoder(realnet, 340)
+def model_load_weights(identifier: str):
+    # Download weights (Brain-Score team modification)
+    # Read the version id and sha1 from json file called "weights.json"
+    weights_info = load_config("weights.json")
+
+    version_id = weights_info['version_ids'][identifier]
+    sha1 = weights_info['sha1s'][identifier]
+
+    weights_path = load_weight_file(bucket="brainscore-storage", folder_name="brainscore-vision/models",
+                                        relative_path=f"ReAlnet/{identifier}_best_model_params.pt",
+                                        version_id=version_id,
+                                        sha1=sha1)
+
+    # Load weights onto CPU and remove "module." from keys
+    weights = torch.load(weights_path, map_location='cpu')
+    new_state_dict = {}
+    for key, val in weights.items():
+        # remove "module." (if it exists) from the key
+        new_key = key.replace("module.", "")
+        new_state_dict[new_key] = val
+
+    encoder.load_state_dict(new_state_dict)
+
+    # Retrieve the realnet portion from the encoder
+    realnet = encoder.realnet
+    realnet.eval()
+    return realnet
+
+def get_model(identifier: str):
+    model = model_load_weights(identifier)
+    preprocessing = functools.partial(load_preprocess_images, image_size=224)
+    wrapper = PytorchWrapper(identifier=identifier, model=model, preprocessing=preprocessing)
+    wrapper.image_size = 224
+    return wrapper
+
+# if __name__ == "__main__":

brainscore_vision/models/ReAlnet/requirements.txt

@@ -0,0 +1,7 @@
+
+torch
+torchvision
+pandas
+numpy
+h5py
+gdown

brainscore_vision/models/ReAlnet/weights.json

@@ -0,0 +1,26 @@
+{
+    "version_ids": {
+      "ReAlnet01": "75oY3CnI17U5S1f_yrZxl1XGhRfJEG9N",
+      "ReAlnet02": "TfGdm1CphJJ1vvkJGcm3n266PHvTuOaV",
+      "ReAlnet03": "dmohrH_AHZzgL_o8Xd2SDp6XCnjPOdAu",
+      "ReAlnet04": "45qJFXHihmIHdpHbjKWZco6STH1eh49p",
+      "ReAlnet05": "nqvoYgiBTyWSskjnpF9YOK4yYQfOnc_H",
+      "ReAlnet06": "6.cloFvnMihiicwQ0jkag8reEe4bVlxZ",
+      "ReAlnet07": "WKJaiN4b1ttpbGYNn8yVjng4LjCqWdk.",
+      "ReAlnet08": "vmouew6ePkPnKP.We8VnVxU7TifuhL.x",
+      "ReAlnet09": "53gqQ2tgS.5MEoncipy9mrBEqCc5izw5",
+      "ReAlnet10": "ZZFMhTm9KQYEXl8OGwKmnTr0S.pxkU0J"
+    },
+    "sha1s": {
+      "ReAlnet01": "05e4e401e8734b97e561aad306fc584b7e027225",
+      "ReAlnet02": "e85769fadb3c09ff88a7d73b01451b6bcccefd77",
+      "ReAlnet03": "f32d01d73380374ae501a1504e9c8cd219e9f0bf",
+      "ReAlnet04": "8062373fd6a74c52360420619235590d3688b4df",
+      "ReAlnet05": "88ca110f6b6d225b7b4e7dca02d2e7a906f5a8ed",
+      "ReAlnet06": "a1658c15a3c9d61262f87349c9fb7aa63854ac5b",
+      "ReAlnet07": "6a1c260839c75f6e6c018e06830562cdcda877e5",
+      "ReAlnet08": "1772211b27dd3a7d9255ac59d5f9b7e7cb6c3314",
+      "ReAlnet09": "159d96f0433a87c7063259dac4527325a3c7b79a",
+      "ReAlnet10": "dbdeaee9280267613ebce92dd5d515d89b544352"
+    }
+  }

brainscore_vision/models/ReAlnet_cornet/__init__.py

@@ -0,0 +1,46 @@
+from brainscore_vision import model_registry
+from .helpers.helpers import CORnetCommitment, _build_time_mappings
+from .model import get_model, TIME_MAPPINGS, get_layers
+
+
+model_registry['ReAlnet01_cornet'] = lambda: CORnetCommitment(identifier='ReAlnet01_cornet', activations_model=get_model('ReAlnet01_cornet'),
+                                                              layers=get_layers('ReAlnet01_cornet'),
+                                                      time_mapping=_build_time_mappings(TIME_MAPPINGS))
+
+
+model_registry['ReAlnet02_cornet'] = lambda: CORnetCommitment(identifier='ReAlnet02_cornet', activations_model=get_model('ReAlnet02_cornet'),
+                                                              layers=get_layers('ReAlnet02_cornet'),
+                                                        time_mapping=_build_time_mappings(TIME_MAPPINGS))
+
+model_registry['ReAlnet03_cornet'] = lambda: CORnetCommitment(identifier='ReAlnet03_cornet', activations_model=get_model('ReAlnet03_cornet'),
+                                                              layers = get_layers('ReAlnet03_cornet'),
+                                                                time_mapping=_build_time_mappings(TIME_MAPPINGS))
+
+model_registry['ReAlnet04_cornet'] = lambda: CORnetCommitment(identifier='ReAlnet04_cornet', activations_model=get_model('ReAlnet04_cornet'),
+                                                              layers = get_layers('ReAlnet04_cornet'),
+                                                                time_mapping=_build_time_mappings(TIME_MAPPINGS))
+
+model_registry['ReAlnet05_cornet'] = lambda: CORnetCommitment(identifier='ReAlnet05_cornet', activations_model=get_model('ReAlnet05_cornet'),
+                                                              layers = get_layers('ReAlnet05_cornet'),
+                                                                time_mapping=_build_time_mappings(TIME_MAPPINGS))
+
+model_registry['ReAlnet06_cornet'] = lambda: CORnetCommitment(identifier='ReAlnet06_cornet', activations_model=get_model('ReAlnet06_cornet'),
+                                                              layers = get_layers('ReAlnet06_cornet'),
+                                                                time_mapping=_build_time_mappings(TIME_MAPPINGS))
+
+model_registry['ReAlnet07_cornet'] = lambda: CORnetCommitment(identifier='ReAlnet07_cornet', activations_model=get_model('ReAlnet07_cornet'),
+                                                              layers = get_layers('ReAlnet07_cornet'),
+                                                                time_mapping=_build_time_mappings(TIME_MAPPINGS))
+
+model_registry['ReAlnet08_cornet'] = lambda: CORnetCommitment(identifier='ReAlnet08_cornet', activations_model=get_model('ReAlnet08_cornet'),
+                                                              layers = get_layers('ReAlnet08_cornet'),
+                                                                time_mapping=_build_time_mappings(TIME_MAPPINGS))
+
+model_registry['ReAlnet09_cornet'] = lambda: CORnetCommitment(identifier='ReAlnet09_cornet', activations_model=get_model('ReAlnet09_cornet'),
+                                                              layers = get_layers('ReAlnet09_cornet'),
+                                                                time_mapping=_build_time_mappings(TIME_MAPPINGS))
+
+model_registry['ReAlnet10_cornet'] = lambda: CORnetCommitment(identifier='ReAlnet10_cornet', activations_model=get_model('ReAlnet10_cornet'),
+                                                              layers = get_layers('ReAlnet10_cornet'),
+                                                                time_mapping=_build_time_mappings(TIME_MAPPINGS))
+

brainscore_vision/models/ReAlnet_cornet/helpers/helpers.py

@@ -0,0 +1,215 @@
+import re
+from collections import defaultdict
+from typing import Dict, Tuple
+
+import numpy as np
+from tqdm import tqdm
+
+from brainio.assemblies import merge_data_arrays, NeuroidAssembly, walk_coords
+from brainscore_vision.model_helpers.activations.core import ActivationsExtractorHelper
+from brainscore_vision.model_helpers.activations.pytorch import PytorchWrapper
+from brainscore_vision.model_helpers.brain_transformation.behavior import BehaviorArbiter, LogitsBehavior, \
+    ProbabilitiesMapping, OddOneOut
+from brainscore_vision.model_interface import BrainModel
+from result_caching import store
+
+
+class TemporalPytorchWrapper(PytorchWrapper):
+    def __init__(self, *args, separate_time=True, **kwargs):
+        self._separate_time = separate_time
+        super(TemporalPytorchWrapper, self).__init__(*args, **kwargs)
+
+    def _build_extractor(self, *args, **kwargs):
+        if self._separate_time:
+            return TemporalExtractor(*args, **kwargs)
+        else:
+            return super(TemporalPytorchWrapper, self)._build_extractor(*args, **kwargs)
+
+    def get_activations(self, images, layer_names):
+        # reset
+        self._layer_counter = defaultdict(lambda: 0)
+        self._layer_hooks = {}
+        return super(TemporalPytorchWrapper, self).get_activations(images=images, layer_names=layer_names)
+
+    def register_hook(self, layer, layer_name, target_dict):
+        layer_name = self._strip_layer_timestep(layer_name)
+        if layer_name in self._layer_hooks:  # add hook only once for multiple timesteps
+            return self._layer_hooks[layer_name]
+
+        def hook_function(_layer, _input, output):
+            target_dict[f"{layer_name}-t{self._layer_counter[layer_name]}"] = PytorchWrapper._tensor_to_numpy(output)
+            self._layer_counter[layer_name] += 1
+
+        hook = layer.register_forward_hook(hook_function)
+        self._layer_hooks[layer_name] = hook
+        return hook
+
+    def get_layer(self, layer_name):
+        layer_name = self._strip_layer_timestep(layer_name)
+        return super(TemporalPytorchWrapper, self).get_layer(layer_name)
+
+    def _strip_layer_timestep(self, layer_name):
+        match = re.search('-t[0-9]+$', layer_name)
+        if match:
+            layer_name = layer_name[:match.start()]
+        return layer_name
+
+
+class CORnetCommitment(BrainModel):
+    """
+    CORnet commitment where only the model interface is implemented and behavioral readouts are attached.
+    Importantly, layer-region commitments do not occur due to the anatomical pre-mapping.
+    Further, due to the temporal component of the model, requested time-bins are matched to the nearest committed
+    time-bin for the model.
+    """
+
+    def __init__(self, identifier, activations_model, layers,
+                 time_mapping: Dict[str, Dict[int, Tuple[int, int]]], behavioral_readout_layer=None,
+                 visual_degrees=8):
+        """
+        :param time_mapping: mapping from region -> {model_timestep -> (time_bin_start, time_bin_end)}
+        """
+        self.layers = layers
+        self.activations_model = activations_model
+        self.time_mapping = time_mapping
+        self.recording_layers = None
+        self.recording_time_bins = None
+        self._identifier = identifier
+
+        logits_behavior = LogitsBehavior(
+            identifier=identifier, activations_model=TemporalIgnore(activations_model))
+        behavioral_readout_layer = behavioral_readout_layer or layers[-1]
+        probabilities_behavior = ProbabilitiesMapping(
+            identifier=identifier, activations_model=TemporalIgnore(activations_model), layer=behavioral_readout_layer)
+        odd_one_out = OddOneOut(identifier=identifier, activations_model=TemporalIgnore(activations_model),
+                                layer=behavioral_readout_layer)
+        self.behavior_model = BehaviorArbiter({BrainModel.Task.label: logits_behavior,
+                                               BrainModel.Task.probabilities: probabilities_behavior,
+                                               BrainModel.Task.odd_one_out: odd_one_out,
+                                               })
+        self.do_behavior = False
+
+        self._visual_degrees = visual_degrees
+
+    @property
+    def identifier(self):
+        return self._identifier
+
+    def visual_degrees(self) -> int:
+        return self._visual_degrees
+
+    def start_recording(self, recording_target, time_bins):
+        self.recording_layers = [layer for layer in self.layers if layer.startswith(recording_target)]
+        self.recording_time_bins = time_bins
+
+    def start_task(self, task: BrainModel.Task, *args, **kwargs):
+        if task != BrainModel.Task.passive:
+            self.behavior_model.start_task(task, *args, **kwargs)
+            self.do_behavior = True
+
+    def look_at(self, stimuli, number_of_trials: int = 1, require_variance: bool = False):
+        if self.do_behavior:
+            return self.behavior_model.look_at(stimuli,
+                                               number_of_trials=number_of_trials, require_variance=require_variance)
+        else:
+            # cache, since piecing times together is not too fast unfortunately
+            return self.look_at_cached(self.identifier, stimuli.identifier, stimuli,
+                                       number_of_trials=number_of_trials, require_variance=require_variance)
+
+    @store(identifier_ignore=['stimuli', 'number_of_trials', 'require_variance'])
+    def look_at_cached(self, model_identifier, stimuli_identifier, stimuli,
+                       number_of_trials, require_variance):
+        responses = self.activations_model(stimuli, layers=self.recording_layers,
+                                           number_of_trials=number_of_trials, require_variance=require_variance)
+        # map time
+        regions = set(responses['region'].values)
+        if len(regions) > 1:
+            raise NotImplementedError("cannot handle more than one simultaneous region")
+        region = list(regions)[0]
+        time_bins = [self.time_mapping[region][timestep] if timestep in self.time_mapping[region] else (None, None)
+                     for timestep in responses['time_step'].values]
+        responses['time_bin_start'] = 'time_step', [time_bin[0] for time_bin in time_bins]
+        responses['time_bin_end'] = 'time_step', [time_bin[1] for time_bin in time_bins]
+        responses = NeuroidAssembly(responses.rename({'time_step': 'time_bin'}))
+        responses = responses[{'time_bin': [not np.isnan(time_start) for time_start in responses['time_bin_start']]}]
+        # select time
+        time_responses = []
+        for time_bin in tqdm(self.recording_time_bins, desc='CORnet-time to recording time'):
+            time_bin = time_bin if not isinstance(time_bin, np.ndarray) else time_bin.tolist()
+            time_bin_start, time_bin_end = time_bin
+            nearest_start = find_nearest(responses['time_bin_start'].values, time_bin_start)
+            bin_responses = responses.sel(time_bin_start=nearest_start)
+            bin_responses = NeuroidAssembly(bin_responses.values, coords={
+                **{coord: (dims, values) for coord, dims, values in walk_coords(bin_responses)
+                   if coord not in ['time_bin_level_0', 'time_bin_end']},
+                **{'time_bin_start': ('time_bin', [time_bin_start]),
+                   'time_bin_end': ('time_bin', [time_bin_end])}
+            }, dims=bin_responses.dims)
+            time_responses.append(bin_responses)
+        responses = merge_data_arrays(time_responses)
+        return responses
+
+
+def find_nearest(array, value):
+    array = np.asarray(array)
+    idx = (np.abs(array - value)).argmin()
+    return array[idx]
+
+
+class TemporalIgnore:
+    """
+    Wrapper around a activations model that squeezes out the temporal axis.
+    Useful when there is only one time step and the behavioral readout does not know what to do with time.
+    """
+
+    def __init__(self, temporal_activations_model):
+        self._activations_model = temporal_activations_model
+
+    def __call__(self, *args, **kwargs):
+        activations = self._activations_model(*args, **kwargs)
+        activations = activations.squeeze('time_step')
+        return activations
+
+
+class TemporalExtractor(ActivationsExtractorHelper):
+    # `from_paths` is the earliest method at which we can interject because calls below are stored and checked for the
+    # presence of all layers which, for CORnet, are passed as e.g. `IT.output-t0`.
+    # This code re-arranges the time component.
+    def from_paths(self, *args, **kwargs):
+        raw_activations = super(TemporalExtractor, self).from_paths(*args, **kwargs)
+        # introduce time dimension
+        regions = defaultdict(list)
+        for layer in set(raw_activations['layer'].values):
+            match = re.match(r'(([^-]*)\..*|logits|avgpool)-t([0-9]+)', layer)
+            region, timestep = match.group(2) if match.group(2) else match.group(1), match.group(3)
+            stripped_layer = match.group(1)
+            regions[region].append((layer, stripped_layer, timestep))
+        activations = {}
+        for region, time_layers in regions.items():
+            for (full_layer, stripped_layer, timestep) in time_layers:
+                region_time_activations = raw_activations.sel(layer=full_layer)
+                region_time_activations['layer'] = 'neuroid', [stripped_layer] * len(region_time_activations['neuroid'])
+                activations[(region, timestep)] = region_time_activations
+        for key, key_activations in activations.items():
+            region, timestep = key
+            key_activations['region'] = 'neuroid', [region] * len(key_activations['neuroid'])
+            activations[key] = NeuroidAssembly([key_activations.values], coords={
+                **{coord: (dims, values) for coord, dims, values in walk_coords(activations[key])
+                   if coord != 'neuroid_id'},  # otherwise, neuroid dim will be as large as before with nans
+                **{'time_step': [int(timestep)]}
+            }, dims=['time_step'] + list(key_activations.dims))
+        activations = list(activations.values())
+        activations = merge_data_arrays(activations)
+        # rebuild neuroid_id without timestep
+        neuroid_id = [".".join([f"{value}" for value in values]) for values in zip(*[
+            activations[coord].values for coord in ['model', 'region', 'neuroid_num']])]
+        activations['neuroid_id'] = 'neuroid', neuroid_id
+        return activations
+
+
+def _build_time_mappings(time_mappings):
+    return {region: {
+        timestep: (time_start + timestep * time_step_size,
+                   time_start + (timestep + 1) * time_step_size)
+        for timestep in range(0, timesteps)}
+        for region, (time_start, time_step_size, timesteps) in time_mappings.items()}

brainscore_vision/models/ReAlnet_cornet/model.py

@@ -0,0 +1,69 @@
+import functools
+import importlib
+from brainscore_vision.model_helpers.activations.pytorch import load_preprocess_images
+from brainscore_vision.model_helpers.check_submission import check_models
+import torch.hub
+import ssl
+from brainscore_vision.model_helpers.s3 import load_weight_file
+from torch.nn import Module
+from .helpers.helpers import TemporalPytorchWrapper
+from pathlib import Path
+from urllib.request import urlretrieve
+
+ssl._create_default_https_context = ssl._create_unverified_context
+
+
+TIME_MAPPINGS = {
+        'V1': (50, 100, 1),
+        'V2': (70, 100, 2),
+        # 'V2': (20, 50, 2),  # MS: This follows from the movshon anesthesized-monkey recordings, so might not hold up
+        'V4': (90, 50, 4),
+        'IT': (100, 100, 2),
+    }
+
+
+def get_model(identifier: str):
+
+    class Wrapper(Module):
+        def __init__(self, model):
+            super(Wrapper, self).__init__()
+            self.module = model
+
+    mod = importlib.import_module(f'cornet.cornet_s')
+    model_ctr = getattr(mod, 'CORnet_S')
+    model = model_ctr()
+    model = Wrapper(model)  # model was wrapped with DataParallel, so weights require `module.` prefix
+    # cornet version: shorten identifier
+    identifier_short = identifier[:9]
+    # cornet version: shorten identifier
+    identifier_short = identifier[:9]
+    url = f'https://brainscore-storage.s3.us-east-2.amazonaws.com/brainscore-vision/models/ReAlnet/{identifier_short}_best_model_params.pt'
+    fh = urlretrieve(url, f'{identifier_short}_best_model_params.pth')
+    load_path = fh[0]
+    checkpoint = torch.load(load_path, map_location=lambda storage, loc: storage)  # map onto cpu
+    new_state_dict = {}
+    for key, val in checkpoint.items():
+        # remove "module." (if it exists) from the key
+        new_key = key.replace("realnet.", "")
+        # discard the keys starting with "fc" 
+        if not new_key.startswith('fc'):
+            new_state_dict[new_key] = val
+
+    model.load_state_dict(new_state_dict)
+    model = model.module  # unwrap
+    preprocessing = functools.partial(load_preprocess_images, image_size=224)
+    wrapper = TemporalPytorchWrapper(identifier=identifier, model=model, preprocessing=preprocessing,
+                                     separate_time=True)
+    wrapper.image_size = 224
+    return wrapper
+
+
+
+
+def get_layers(identifier: str):
+    return (['V1.output-t0'] +
+               [f'{area}.output-t{timestep}'
+                for area, timesteps in [('V2', range(2)), ('V4', range(4)), ('IT', range(2))]
+                for timestep in timesteps] +
+               ['decoder.avgpool-t0']
+            )

brainscore_vision/models/ReAlnet_cornet/requirements.txt

@@ -0,0 +1,8 @@
+
+CORnet @ git+https://github.com/dicarlolab/CORnet.git
+torch
+torchvision
+pandas
+numpy
+h5py
+gdown

brainscore_vision/models/Res2Net50_26w_4s/__init__.py

@@ -0,0 +1,5 @@
+from brainscore_vision import model_registry
+from brainscore_vision.model_helpers.brain_transformation import ModelCommitment
+from .model import get_model, get_layers
+
+model_registry['Res2Net50_26w_4s'] = lambda: ModelCommitment(identifier='Res2Net50_26w_4s', activations_model=get_model('Res2Net50_26w_4s'), layers=get_layers('Res2Net50_26w_4s'))