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

brainscore_vision/models/Res2Net50_26w_4s/helpers/resnet_helpers.py

@@ -0,0 +1,161 @@
+from brainscore_vision.model_helpers.s3 import load_weight_file
+import torch
+import torch.nn as nn
+import math
+
+
+class Bottle2neck(nn.Module):
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, baseWidth=26, scale=4, stype='normal'):
+        """ Constructor
+        Args:
+            inplanes: input channel dimensionality
+            planes: output channel dimensionality
+            stride: conv stride. Replaces pooling layer.
+            downsample: None when stride = 1
+            baseWidth: basic width of conv3x3
+            scale: number of scale.
+            type: 'normal': normal set. 'stage': first block of a new stage.
+        """
+        super(Bottle2neck, self).__init__()
+
+        width = int(math.floor(planes * (baseWidth / 64.0)))
+        self.conv1 = nn.Conv2d(inplanes, width * scale, kernel_size=1, bias=False)
+        self.bn1 = nn.BatchNorm2d(width * scale)
+
+        if scale == 1:
+            self.nums = 1
+        else:
+            self.nums = scale - 1
+        if stype == 'stage':
+            self.pool = nn.AvgPool2d(kernel_size=3, stride=stride, padding=1)
+        convs = []
+        bns = []
+        for i in range(self.nums):
+            convs.append(nn.Conv2d(width, width, kernel_size=3, stride=stride, padding=1, bias=False))
+            bns.append(nn.BatchNorm2d(width))
+        self.convs = nn.ModuleList(convs)
+        self.bns = nn.ModuleList(bns)
+
+        self.conv3 = nn.Conv2d(width * scale, planes * self.expansion, kernel_size=1, bias=False)
+        self.bn3 = nn.BatchNorm2d(planes * self.expansion)
+
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stype = stype
+        self.scale = scale
+        self.width = width
+
+    def forward(self, x):
+        residual = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        spx = torch.split(out, self.width, 1)
+        for i in range(self.nums):
+            if i == 0 or self.stype == 'stage':
+                sp = spx[i]
+            else:
+                sp = sp + spx[i]
+            sp = self.convs[i](sp)
+            sp = self.relu(self.bns[i](sp))
+            if i == 0:
+                out = sp
+            else:
+                out = torch.cat((out, sp), 1)
+        if self.scale != 1 and self.stype == 'normal':
+            out = torch.cat((out, spx[self.nums]), 1)
+        elif self.scale != 1 and self.stype == 'stage':
+            out = torch.cat((out, self.pool(spx[self.nums])), 1)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            residual = self.downsample(x)
+
+        out += residual
+        out = self.relu(out)
+
+        return out
+
+
+class Res2Net(nn.Module):
+
+    def __init__(self, block, layers, baseWidth=26, scale=4, num_classes=1000):
+        self.inplanes = 64
+        super(Res2Net, self).__init__()
+        self.baseWidth = baseWidth
+        self.scale = scale
+        self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3,
+                               bias=False)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0])
+        self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+        self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+        self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
+        self.avgpool = nn.AdaptiveAvgPool2d(1)
+        self.fc = nn.Linear(512 * block.expansion, num_classes)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, nn.BatchNorm2d):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+    def _make_layer(self, block, planes, blocks, stride=1):
+        downsample = None
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                nn.Conv2d(self.inplanes, planes * block.expansion,
+                          kernel_size=1, stride=stride, bias=False),
+                nn.BatchNorm2d(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample=downsample,
+                            stype='stage', baseWidth=self.baseWidth, scale=self.scale))
+        self.inplanes = planes * block.expansion
+        for i in range(1, blocks):
+            layers.append(block(self.inplanes, planes, baseWidth=self.baseWidth, scale=self.scale))
+
+        return nn.Sequential(*layers)
+
+    def forward(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        x = self.avgpool(x)
+        x = x.view(x.size(0), -1)
+        x = self.fc(x)
+
+        return x
+
+
+def res2net50_26w_4s(pretrained=False, **kwargs):
+    """Constructs a Res2Net-50_26w_4s model.
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+    """
+    device = "cpu"
+    model = Res2Net(Bottle2neck, [3, 4, 6, 3], baseWidth=26, scale=4, **kwargs)
+    weights_path = load_weight_file(bucket="brainscore-storage", folder_name="brainscore-vision/models",
+                                    relative_path="Res2Net50_26w_4s/res2net50_26w_4s-06e79181.pth",
+                                    version_id="null",
+                                    sha1="905ad1183c7c8fb4b13f0df8c1a22bb76ee2a2fb")
+    checkpoint = torch.load(weights_path, map_location=device)
+    model.load_state_dict(checkpoint)
+    return model

brainscore_vision/models/Res2Net50_26w_4s/model.py

@@ -0,0 +1,75 @@
+from brainscore_vision.model_helpers.check_submission import check_models
+from .helpers.resnet_helpers import res2net50_26w_4s
+import functools
+from brainscore_vision.model_helpers.activations.pytorch import PytorchWrapper
+from brainscore_vision.model_helpers.activations.pytorch import load_preprocess_images
+
+
+
+def get_model(name):
+    """
+    This method fetches an instance of a base model. The instance has to be callable and return a xarray object,
+    containing activations. There exist standard wrapper implementations for common libraries, like pytorch and
+    keras. Checkout the examples folder, to see more. For custom implementations check out the implementation of the
+    wrappers.
+    :param name: the name of the model to fetch
+    :return: the model instance
+    """
+
+    assert name == 'Res2Net50_26w_4s'
+
+    model = res2net50_26w_4s(pretrained=False)
+    preprocessing = functools.partial(load_preprocess_images, image_size=224)
+    wrapper = PytorchWrapper(identifier='Res2Net50_26w_4s', model=model, preprocessing=preprocessing)
+    wrapper.image_size = 224
+    return wrapper
+
+
+# get_layers method to tell the code what layers to consider. If you are submitting a custom
+# model, then you will most likley need to change this method's return values.
+def get_layers(name):
+    """
+    This method returns a list of string layer names to consider per model. The benchmarks maps brain regions to
+    layers and uses this list as a set of possible layers. The lists doesn't have to contain all layers, the less the
+    faster the benchmark process works. Additionally the given layers have to produce an activations vector of at least
+    size 25! The layer names are delivered back to the model instance and have to be resolved in there. For a pytorch
+    model, the layer name are for instance dot concatenated per module, e.g. "features.2".
+    :param name: the name of the model, to return the layers for
+    :return: a list of strings containing all layers, that should be considered as brain area.
+    """
+
+    assert name == 'Res2Net50_26w_4s'
+    return ['maxpool', 'layer1.0', 'layer1.1', 'layer1.2',
+            'layer2.0', 'layer2.1', 'layer2.2', 'layer2.3',
+            'layer3.0', 'layer3.1', 'layer3.2', 'layer3.3',
+            'layer3.4', 'layer3.5', 'layer4.0', 'layer4.1', 'layer4.2', 'avgpool', 'fc']
+
+
+def get_bibtex(model_identifier):
+    """
+    A method returning the bibtex reference of the requested model as a string.
+    """
+
+    # from pytorch.py:
+    return '''@article{DBLP:journals/corr/abs-1904-01169,
+  author       = {Shanghua Gao and
+                  Ming{-}Ming Cheng and
+                  Kai Zhao and
+                  Xinyu Zhang and
+                  Ming{-}Hsuan Yang and
+                  Philip H. S. Torr},
+  title        = {Res2Net: {A} New Multi-scale Backbone Architecture},
+  journal      = {CoRR},
+  volume       = {abs/1904.01169},
+  year         = {2019},
+  url          = {http://arxiv.org/abs/1904.01169},
+  eprinttype    = {arXiv},
+  eprint       = {1904.01169},
+  timestamp    = {Thu, 25 Apr 2019 10:24:54 +0200},
+  biburl       = {https://dblp.org/rec/journals/corr/abs-1904-01169.bib},
+  bibsource    = {dblp computer science bibliography, https://dblp.org}
+}'''
+
+# Main Method: In submitting a custom model, you should not have to mess with this.
+if __name__ == '__main__':
+    check_models.check_base_models(__name__)

brainscore_vision/models/Res2Net50_26w_4s/region_layer_map/Res2Net50_26w_4s.json

@@ -0,0 +1 @@
+{"IT": "layer3.5", "V4": "layer2.1", "V1": "layer1.0", "V2": "layer3.0"}

brainscore_vision/models/Res2Net50_26w_4s/requirements.txt

@@ -0,0 +1 @@
+torch

brainscore_vision/models/Res2Net50_26w_4s/test.py

@@ -0,0 +1,8 @@
+import brainscore_vision
+import pytest
+
+
+@pytest.mark.travis_slow
+def test_has_identifier():
+    model = brainscore_vision.load_model('Res2Net50_26w_4s')
+    assert model.identifier == 'Res2Net50_26w_4s'

brainscore_vision/models/VOneCORnet_S/__init__.py

@@ -0,0 +1,9 @@
+from brainscore_vision.model_helpers.brain_transformation import ModelCommitment
+from brainscore_vision import model_registry
+from .model import get_layers,get_model
+from .helpers.vonecornets import VOneCORnetCommitment, _build_time_mappings, CORNET_S_TIMEMAPPING
+
+
+model_registry['VOneCORnet-S'] = \
+    lambda: VOneCORnetCommitment(identifier='VOneCORnet-S', activations_model=get_model('VOneCORnet-S'),
+                                layers=get_layers('VOneCORnet-S'), time_mapping=_build_time_mappings(CORNET_S_TIMEMAPPING))

brainscore_vision/models/VOneCORnet_S/helpers/cornet_helpers.py

@@ -0,0 +1,34 @@
+import logging
+from importlib import import_module
+from brainscore_vision.model_helpers.activations import PytorchWrapper
+from brainscore_vision.model_helpers.activations.pytorch import load_preprocess_images
+from .cornet_s_helpers import cornet as cornet_ctr
+from vonenet import get_model
+from .cornet_s_helpers import TemporalPytorchWrapper
+import functools
+
+
+_logger = logging.getLogger(__name__)
+
+
+def torchvision_model(identifier, image_size):
+    module = import_module(f'torchvision.models')
+    model_ctr = getattr(module, identifier)
+    from model_helpers.activations.pytorch import load_preprocess_images
+    preprocessing = functools.partial(load_preprocess_images, image_size=image_size)
+    wrapper = PytorchWrapper(identifier=identifier, model=model_ctr(pretrained=True), preprocessing=preprocessing)
+    wrapper.image_size = image_size
+    return wrapper
+
+def cornet(*args, **kwargs):  # wrapper to avoid having to import cornet at top-level
+    return cornet_ctr(*args, **kwargs)
+
+def vonecornet(model_name='cornets'):
+    model = get_model(model_name)
+    model = model.module
+    preprocessing = functools.partial(load_preprocess_images, image_size=224,
+                                      normalize_mean=(0.5, 0.5, 0.5), normalize_std=(0.5, 0.5, 0.5))
+    wrapper = TemporalPytorchWrapper(identifier='vone'+model_name, model=model, preprocessing=preprocessing)
+    wrapper.image_size = 224
+    return wrapper
+

brainscore_vision/models/VOneCORnet_S/helpers/cornet_s_helpers.py

@@ -0,0 +1,128 @@
+import importlib
+import logging
+from collections import defaultdict
+import functools
+import re
+import torch
+from torch.nn import Module
+from brainio.assemblies import merge_data_arrays, NeuroidAssembly, walk_coords
+from brainscore_vision.model_helpers.s3 import load_weight_file
+from brainscore_vision.model_helpers.activations.core import ActivationsExtractorHelper
+from brainscore_vision.model_helpers.activations.pytorch import PytorchWrapper
+from brainscore_vision.model_helpers.activations.pytorch import load_preprocess_images
+
+_logger = logging.getLogger(__name__)
+
+
+def cornet(identifier, separate_time=True):
+    cornet_type = re.match('CORnet-(.*)', identifier).group(1)
+    mod = importlib.import_module(f'cornet.cornet_{cornet_type.lower()}')
+    model_ctr = getattr(mod, f'CORnet_{cornet_type.upper()}')
+    model = model_ctr()
+
+    WEIGHT_MAPPING = {
+        'Z': 'cornet_z-5c427c9c.pth',
+        'R': 'cornet_r_epoch25.pth.tar',
+        'S': 'cornet_s_epoch43.pth.tar',
+        'R2': 'cornet_r2_epoch_60.pth.tar',
+        'S10': 'cornet_s10_latest.pth.tar',
+        'S222': 'cornet_s222_latest.pth.tar',
+        'S444': 'cornet_s444_latest.pth.tar',
+        'S484': 'cornet_s484_latest.pth.tar',
+    }
+
+    class Wrapper(Module):
+        def __init__(self, model):
+            super(Wrapper, self).__init__()
+            self.module = model
+
+    model = Wrapper(model)  # model was wrapped with DataParallel, so weights require `module.` prefix
+    weights_path = load_weight_file(bucket="brainscore-storage", folder_name="brainscore-vision/models",
+                                    relative_path="cornet_s/cornet_s_epoch43.pth.tar",
+                                    version_id="null",
+                                    sha1="a4bfd8eda33b45fd945da1b972ab0b7cad38d60f")
+    checkpoint = torch.load(weights_path, map_location=lambda storage, loc: storage)  # map onto cpu
+    model.load_state_dict(checkpoint['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=separate_time)
+    wrapper.image_size = 224
+    return wrapper
+
+
+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 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

brainscore_vision/models/VOneCORnet_S/helpers/cornets.py

@@ -0,0 +1,136 @@
+import logging
+import numpy as np
+from brainio.assemblies import merge_data_arrays, NeuroidAssembly, walk_coords
+from brainscore_vision.model_helpers.brain_transformation import ModelCommitment
+from brainscore_vision.model_helpers.brain_transformation.temporal import fix_timebin_naming
+from tqdm import tqdm
+from typing import Dict, Tuple
+from .cornet_helpers import cornet
+
+_logger = logging.getLogger(__name__)
+
+CORNET_S_TIMEMAPPING = {
+    'V1': (50, 100, 1),
+    'V2': (70, 100, 2),
+    'V4': (90, 50, 4),
+    'IT': (100, 100, 2),
+}
+
+
+class CORnetCommitment(ModelCommitment):
+    """
+    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, *args, time_mapping: Dict[str, Dict[int, Tuple[int, int]]], **kwargs):
+        """
+        :param time_mapping: mapping from region -> {model_timestep -> (time_bin_start, time_bin_end)}
+        """
+        super(CORnetCommitment, self).__init__(*args, **kwargs)
+        self.time_mapping = time_mapping
+        self.recording_layers = None
+        self.recording_time_bins = None
+        # deal with activations_model returning a time_bin
+        for key, executor in self.behavior_model.mapping.items():
+            executor.activations_model = TemporalIgnore(executor.activations_model)
+
+    def start_recording(self, recording_target, time_bins):
+        self.recording_target = recording_target
+        self.recording_layers = [layer for layer in self.layers if layer.startswith(recording_target)]
+        self.recording_time_bins = time_bins
+
+    def look_at(self, stimuli, number_of_trials=1):
+        if self.do_behavior:
+            return super(CORnetCommitment, self).look_at(stimuli, number_of_trials=number_of_trials)
+        else:
+            return self.look_at_temporal(stimuli=stimuli)  # ignore number_of_trials
+
+    def look_at_temporal(self, stimuli):
+        responses = self.activations_model(stimuli, layers=self.recording_layers)
+        # map time
+        if hasattr(self, 'recording_target'):
+            regions = set([self.recording_target])
+        else:
+            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)
+        responses = fix_timebin_naming(responses)  # work around xarray merge bug introduced in 0.16.2
+        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
+
+
+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()}
+
+
+def cornet_z_brainmodel():
+    return CORnetCommitment(identifier='CORnet-Z', activations_model=cornet('CORnet-Z'),
+                            layers=[f'{region}.output-t0' for region in ['V1', 'V2', 'V4', 'IT']] +
+                                   ['decoder.avgpool-t0'],
+                            time_mapping={
+                                'V1': {0: (50, 150)},
+                                'V2': {0: (70, 170)},
+                                'V4': {0: (90, 190)},
+                                'IT': {0: (100, 200)},
+                            })
+
+
+def cornet_s_brainmodel():
+    # map region -> (time_start, time_step_size, timesteps)
+    time_mappings = CORNET_S_TIMEMAPPING
+    return CORnetCommitment(identifier='CORnet-S', activations_model=cornet('CORnet-S'),
+                            layers=['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'],
+                            time_mapping=_build_time_mappings(time_mappings))

brainscore_vision/models/VOneCORnet_S/helpers/vonecornets.py

@@ -0,0 +1,38 @@
+from .cornets import CORnetCommitment, _build_time_mappings, CORNET_S_TIMEMAPPING
+import logging
+from .cornet_helpers import vonecornet
+_logger = logging.getLogger(__name__)
+
+
+class VOneCORnetCommitment(CORnetCommitment):
+    def start_recording(self, recording_target, time_bins):
+        self.recording_target = recording_target
+        if recording_target == 'V1':
+            self.recording_layers = ['vone_block.output-t0']
+        else:
+            self.recording_layers = [layer for layer in self.layers if recording_target in layer]
+        self.recording_time_bins = time_bins
+
+    def look_at(self, stimuli, number_of_trials=1):
+        stimuli_identifier = stimuli.identifier
+        for trial_number in range(number_of_trials):
+            if stimuli_identifier:
+                stimuli.identifier = stimuli_identifier + '-trial' + f'{trial_number:03d}'
+            if trial_number == 0:
+                activations = super().look_at(stimuli, number_of_trials=1)
+            else:
+                activations += super().look_at(stimuli, number_of_trials=1)
+        stimuli.identifier = stimuli_identifier
+        return activations/number_of_trials
+
+
+def vonecornet_s_brainmodel():
+    # map region -> (time_start, time_step_size, timesteps)
+    time_mappings = CORNET_S_TIMEMAPPING
+    return VOneCORnetCommitment(identifier='VOneCORnet-S', activations_model=vonecornet('cornets'),
+                                layers=['vone_block.output-t0'] + [f'model.{area}.output-t{timestep}'
+                                                                   for area, timesteps in
+                                                                   [('V2', range(2)), ('V4', range(4)), ('IT', range(2))]
+                                                                   for timestep in timesteps] +
+                                       ['model.decoder.avgpool-t0'], time_mapping=_build_time_mappings(time_mappings))
+

brainscore_vision/models/VOneCORnet_S/model.py

@@ -0,0 +1,25 @@
+from brainscore_vision.model_helpers.check_submission import check_models
+import functools
+import torchvision.models
+from brainscore_vision.model_helpers.activations.pytorch import PytorchWrapper
+from brainscore_vision.model_helpers.activations.pytorch import load_preprocess_images
+from .helpers.cornet_helpers import vonecornet
+
+
+def get_model(name):
+    assert name == 'VOneCORnet-S'
+    return vonecornet('cornets')
+
+def get_layers(name):
+    assert name == 'VOneCORnet-S'
+    return ['vone_block.output-t0'] + [f'model.{area}.output-t{timestep}'
+                                                                   for area, timesteps in
+                                                                   [('V2', range(2)), ('V4', range(4)), ('IT', range(2))]
+                                                                   for timestep in timesteps] + ['model.decoder.avgpool-t0']
+
+def get_bibtex(model_identifier):
+    return """xx"""
+
+
+if __name__ == '__main__':
+    check_models.check_base_models(__name__)

brainscore_vision/models/VOneCORnet_S/requirements.txt

@@ -0,0 +1 @@
+vonenet @ git+https://github.com/dicarlolab/vonenet

brainscore_vision/models/VOneCORnet_S/test.py

@@ -0,0 +1,8 @@
+import pytest
+import brainscore_vision
+
+
+@pytest.mark.travis_slow
+def test_has_identifier():
+    model = brainscore_vision.load_model('VOneCORnet-S')
+    assert model.identifier == 'VOneCORnet-S'

brainscore_vision/models/alexnet_training_seed_01/__init__.py

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