brain-score.org submission (user:728) | (public:False)

brainscore_vision/models/alex_hello/__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['alex_hello'] = lambda: ModelCommitment(
+    identifier='alex_hello', activations_model=get_model(), layers=get_layers(),
+)

brainscore_vision/models/alex_hello/model.py

@@ -0,0 +1,9 @@
+import functools, torch
+from torchvision.models import alexnet
+from brainscore_vision.model_helpers.activations.pytorch import PytorchWrapper, load_preprocess_images
+def get_model():
+    m = alexnet(weights='IMAGENET1K_V1'); m.classifier[-1] = torch.nn.Identity()
+    pre = functools.partial(load_preprocess_images, image_size=224)
+    w = PytorchWrapper(identifier='alex_hello', model=m, preprocessing=pre); w.image_size=224
+    return w
+def get_layers(): return ['features.3','features.6','features.8','classifier.5']

brainscore_vision/models/alex_hello/requirements.txt

@@ -0,0 +1,2 @@
+torch
+torchvision

brainscore_vision/models/alex_hello/test.py

@@ -0,0 +1,2 @@
+import brainscore_vision as bsv
+def test_loads(): assert bsv.load_model('alex_hello').identifier == 'alex_hello'