add visual search

.gitignore

@@ -1,2 +1,7 @@
 *.pyc
 *.DS_Store
+__pycache__/
+.ipynb_checkpoints
+build/*
+dist/*
+candidate_models.egg-info/*

candidate_models/base_models/__init__.py

@@ -347,55 +347,55 @@ class BaseModelPool(UniqueKeyDict):
     Each entry maps from `name` to an activations extractor.
     """
 
-    def __init__(self):
+    def __init__(self, input_size=None):
         super(BaseModelPool, self).__init__(reload=True)
 
         _key_functions = {
-            'alexnet': lambda: pytorch_model('alexnet', image_size=224),
-            'squeezenet1_0': lambda: pytorch_model('squeezenet1_0', image_size=224),
-            'squeezenet1_1': lambda: pytorch_model('squeezenet1_1', image_size=224),
-            'resnet-18': lambda: pytorch_model('resnet18', image_size=224),
-            'resnet-34': lambda: pytorch_model('resnet34', image_size=224),
-            'resnet-50-pytorch': lambda: pytorch_model('resnet50', image_size=224),
-            'resnet-50-robust': lambda: robust_model('resnet50', image_size=224),
-
-            'vgg-16': lambda: keras_model('vgg16', 'VGG16', image_size=224),
-            'vgg-19': lambda: keras_model('vgg19', 'VGG19', image_size=224),
+            'alexnet': lambda: pytorch_model('alexnet', image_size=224 if input_size==None else input_size),
+            'squeezenet1_0': lambda: pytorch_model('squeezenet1_0', image_size=224 if input_size==None else input_size),
+            'squeezenet1_1': lambda: pytorch_model('squeezenet1_1', image_size=224 if input_size==None else input_size),
+            'resnet-18': lambda: pytorch_model('resnet18', image_size=224 if input_size==None else input_size),
+            'resnet-34': lambda: pytorch_model('resnet34', image_size=224 if input_size==None else input_size),
+            'resnet-50-pytorch': lambda: pytorch_model('resnet50', image_size=224 if input_size==None else input_size),
+            'resnet-50-robust': lambda: robust_model('resnet50', image_size=224 if input_size==None else input_size),
+
+            'vgg-16': lambda: keras_model('vgg16', 'VGG16', image_size=224 if input_size==None else input_size),
+            'vgg-19': lambda: keras_model('vgg19', 'VGG19', image_size=224 if input_size==None else input_size),
             'vggface': vggface,
-            'xception': lambda: keras_model('xception', 'Xception', image_size=299),
-            'densenet-121': lambda: keras_model('densenet', 'DenseNet121', image_size=224),
-            'densenet-169': lambda: keras_model('densenet', 'DenseNet169', image_size=224),
-            'densenet-201': lambda: keras_model('densenet', 'DenseNet201', image_size=224),
-
-            'inception_v1': lambda: TFSlimModel.init('inception_v1', preprocessing_type='inception', image_size=224),
-            'inception_v2': lambda: TFSlimModel.init('inception_v2', preprocessing_type='inception', image_size=224),
-            'inception_v3': lambda: TFSlimModel.init('inception_v3', preprocessing_type='inception', image_size=299),
-            'inception_v4': lambda: TFSlimModel.init('inception_v4', preprocessing_type='inception', image_size=299),
+            'xception': lambda: keras_model('xception', 'Xception', image_size=299 if input_size==None else input_size),
+            'densenet-121': lambda: keras_model('densenet', 'DenseNet121', image_size=224 if input_size==None else input_size),
+            'densenet-169': lambda: keras_model('densenet', 'DenseNet169', image_size=224 if input_size==None else input_size),
+            'densenet-201': lambda: keras_model('densenet', 'DenseNet201', image_size=224 if input_size==None else input_size),
+
+            'inception_v1': lambda: TFSlimModel.init('inception_v1', preprocessing_type='inception', image_size=224 if input_size==None else input_size),
+            'inception_v2': lambda: TFSlimModel.init('inception_v2', preprocessing_type='inception', image_size=224 if input_size==None else input_size),
+            'inception_v3': lambda: TFSlimModel.init('inception_v3', preprocessing_type='inception', image_size=299 if input_size==None else input_size),
+            'inception_v4': lambda: TFSlimModel.init('inception_v4', preprocessing_type='inception', image_size=299 if input_size==None else input_size),
             'inception_resnet_v2': lambda: TFSlimModel.init('inception_resnet_v2', preprocessing_type='inception',
-                                                            image_size=299),
+                                                            image_size=299 if input_size==None else input_size),
             'resnet-50_v1': lambda: TFSlimModel.init('resnet-50_v1', net_name='resnet_v1_50', preprocessing_type='vgg',
-                                                     image_size=224, labels_offset=0),
+                                                     image_size=224 if input_size==None else input_size, labels_offset=0),
             'resnet-101_v1': lambda: TFSlimModel.init('resnet-101_v1', net_name='resnet_v1_101',
                                                       preprocessing_type='vgg',
-                                                      image_size=224, labels_offset=0),
+                                                      image_size=224 if input_size==None else input_size, labels_offset=0),
             'resnet-152_v1': lambda: TFSlimModel.init('resnet-152_v1', net_name='resnet_v1_152',
                                                       preprocessing_type='vgg',
-                                                      image_size=224, labels_offset=0),
+                                                      image_size=224 if input_size==None else input_size, labels_offset=0),
             # image_size is 299 for resnet-v2, this is a bug in tf-slim.
             # see https://github.com/tensorflow/models/tree/8b18491b26e4b8271db757a3245008882ea112b3/research/slim:
             # "ResNet V2 models use Inception pre-processing and input image size of 299"
             'resnet-50_v2': lambda: TFSlimModel.init('resnet-50_v2', net_name='resnet_v2_50',
                                                      preprocessing_type='inception',
-                                                     image_size=299),
+                                                     image_size=299 if input_size==None else input_size),
             'resnet-101_v2': lambda: TFSlimModel.init('resnet-101_v2', net_name='resnet_v2_101',
                                                       preprocessing_type='inception',
-                                                      image_size=299),
+                                                      image_size=299 if input_size==None else input_size),
             'resnet-152_v2': lambda: TFSlimModel.init('resnet-152_v2', net_name='resnet_v2_152',
                                                       preprocessing_type='inception',
-                                                      image_size=299),
-            'nasnet_mobile': lambda: TFSlimModel.init('nasnet_mobile', preprocessing_type='inception', image_size=331),
-            'nasnet_large': lambda: TFSlimModel.init('nasnet_large', preprocessing_type='inception', image_size=331),
-            'pnasnet_large': lambda: TFSlimModel.init('pnasnet_large', preprocessing_type='inception', image_size=331),
+                                                      image_size=299 if input_size==None else input_size),
+            'nasnet_mobile': lambda: TFSlimModel.init('nasnet_mobile', preprocessing_type='inception', image_size=331 if input_size==None else input_size),
+            'nasnet_large': lambda: TFSlimModel.init('nasnet_large', preprocessing_type='inception', image_size=331 if input_size==None else input_size),
+            'pnasnet_large': lambda: TFSlimModel.init('pnasnet_large', preprocessing_type='inception', image_size=331 if input_size==None else input_size),
             'bagnet9': lambda: bagnet("bagnet9"),
             'bagnet17': lambda: bagnet("bagnet17"),
             'bagnet33': lambda: bagnet("bagnet33"),

candidate_models/model_commitments/__init__.py

@@ -1,7 +1,10 @@
 from candidate_models.base_models import base_model_pool
+from candidate_models.base_models import BaseModelPool
 
 from candidate_models.model_commitments.cornets import cornet_brain_pool
 from candidate_models.model_commitments.model_layer_def import model_layers
+from candidate_models.model_commitments.vs_layer import visual_search_layer
+
 from brainscore.submission.ml_pool import MLBrainPool
 from brainscore.submission.utils import UniqueKeyDict
 
@@ -14,3 +17,18 @@
 
 for identifier, model in cornet_brain_pool.items():
     brain_translated_pool[identifier] = model
+
+def MLSearchPool(target_img_size=28, search_image_size=224):
+    target_model_pool = BaseModelPool(input_size=target_img_size)
+    stimuli_model_pool = BaseModelPool(input_size=search_image_size)
+
+    vs_model_param = {}
+    vs_model_param['tar_pool'] = target_model_pool
+    vs_model_param['stim_pool'] = stimuli_model_pool
+    vs_model_param['model_layers'] = visual_search_layer
+    vs_model_param['tar_size'] = target_img_size
+    vs_model_param['stim_size'] = stimuli_img_size
+
+    ml_search_pool = MLBrainPool(base_model_pool, model_layers, vs_model_param=vs_model_param)
+
+    return ml_search_pool

candidate_models/model_commitments/vs_layer.py

@@ -0,0 +1,9 @@
+import itertools
+
+from brainscore.submission.ml_pool import ModelLayers
+
+layers = {
+    'vgg-16': [f'block{i + 1}_pool' for i in range(3,5)],
+}
+
+visual_search_layer = ModelLayers(layers)

examples/score-model.ipynb

@@ -349,4 +349,4 @@
  },
  "nbformat": 4,
  "nbformat_minor": 1
-}
+}

examples/visual-search-score-model.ipynb

@@ -0,0 +1,122 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import warnings\n",
+    "warnings.simplefilter(\"ignore\")\n",
+    "\n",
+    "try:\n",
+    "    from tensorflow.python.util import module_wrapper as deprecation\n",
+    "except ImportError:\n",
+    "    from tensorflow.python.util import deprecation_wrapper as deprecation\n",
+    "deprecation._PER_MODULE_WARNING_LIMIT = 0"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "WARNING:tensorflow:\n",
+      "The TensorFlow contrib module will not be included in TensorFlow 2.0.\n",
+      "For more information, please see:\n",
+      "  * https://github.com/tensorflow/community/blob/master/rfcs/20180907-contrib-sunset.md\n",
+      "  * https://github.com/tensorflow/addons\n",
+      "  * https://github.com/tensorflow/io (for I/O related ops)\n",
+      "If you depend on functionality not listed there, please file an issue.\n",
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "from candidate_models import score_model\n",
+    "from candidate_models.model_commitments import brain_translated_pool\n",
+    "import numpy as np"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Using TensorFlow backend.\n",
+      "activations: 100%|██████████| 320/320 [00:01<00:00, 231.07it/s]\n",
+      "layer packaging: 100%|██████████| 1/1 [00:00<00:00, 110.52it/s]\n",
+      "activations: 100%|██████████| 320/320 [01:01<00:00,  5.20it/s]\n",
+      "layer packaging: 100%|██████████| 1/1 [00:00<00:00,  2.07it/s]\n",
+      "visual search stimuli: 100%|██████████| 300/300 [01:02<00:00,  4.78it/s]\n",
+      "comparing with human data: 100%|██████████| 15/15 [00:06<00:00,  2.49it/s]"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "<xarray.Score (aggregation: 2)>\n",
+      "array([0.923438, 0.005427])\n",
+      "Coordinates:\n",
+      "  * aggregation  (aggregation) <U6 'center' 'error'\n",
+      "Attributes:\n",
+      "    raw:                   <xarray.Score (aggregation: 2)>\\narray([0.407328, ...\n",
+      "    ceiling:               <xarray.Score (aggregation: 2)>\\narray([0.4411,   ...\n",
+      "    model_identifier:      vgg-16\n",
+      "    benchmark_identifier:  klab.Zhang2018-object_search\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    }
+   ],
+   "source": [
+    "identifier = 'vgg-16'\n",
+    "model = brain_translated_pool[identifier]\n",
+    "score = score_model(model_identifier=identifier, model=model, benchmark_identifier='klab.Zhang2018-object_search')\n",
+    "print(score)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.7.6"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

tests/test_search.py

@@ -0,0 +1,13 @@
+import pytest
+from pytest import approx
+
+from candidate_models.model_commitments import brain_translated_pool
+from brainscore.benchmarks import benchmark_pool
+import numpy as np
+import matplotlib.pyplot as plt
+import brainscore
+
+def test_search():
+    model = brain_translated_pool['vgg-16']
+    score = score_model(model_identifier='vgg-16', model=model, benchmark_identifier='klab.Zhang2018-object_search')
+    assert score.raw.sel(aggregation='center') == approx(0.407328, abs=.005)