added model source parameter; fixes #56

README.md

@@ -9,7 +9,7 @@
 
 ## Model collection
 
-Features can be extracted for all models in [torchvision](https://pytorch.org/vision/0.8/models.html), all models in [Keras](https://www.tensorflow.org/api_docs/python/tf/keras/applications), all models in [timm](https://github.com/rwightman/pytorch-image-models), custom models trained on Ecoset, each of the [CORnet](https://github.com/dicarlolab/CORnet) versions and both [CLIP](https://github.com/openai/CLIP) variants (`clip-ViT` and `clip-RN`). Note, that the respective model name must be used. For example, if you want to use the VGG16 model from torchvision, you will have to use `vgg16` and if you want to use the VGG16 model from Keras, you will have to use the model name `VGG16`.<br>
+Features can be extracted for all models in [torchvision](https://pytorch.org/vision/0.8/models.html), all models in [Keras](https://www.tensorflow.org/api_docs/python/tf/keras/applications), all models in [timm](https://github.com/rwightman/pytorch-image-models), custom models trained on Ecoset, each of the [CORnet](https://github.com/dicarlolab/CORnet) versions and both [CLIP](https://github.com/openai/CLIP) variants (`clip-ViT` and `clip-RN`). Note, that the respective model name must be used. For example, if you want to use the VGG16 model from torchvision, you will have to use `vgg16` and if you want to use the VGG16 model from Keras, you will have to use the model name `VGG16`. You can further specify the model source by setting the `source` parameter.<br>
 For the correct abbreviations of [torchvision](https://pytorch.org/vision/0.8/models.html) models have a look [here](https://github.com/pytorch/vision/tree/master/torchvision/models). For the correct abbreviations of [CORnet](https://github.com/dicarlolab/CORnet) models look [here](https://github.com/dicarlolab/CORnet/tree/master/cornet). To separate the string `cornet` from its variant (e.g., `s`, `z`) use a hyphen instead of an underscore (e.g., `cornet-s`, `cornet-z`).<br>
 
 Examples:  `alexnet`, `resnet18`, `resnet50`, `resnet101`, `vit_b_16`, `vit_b_32`, `vgg13`, `vgg13_bn`, `vgg16`, `vgg16_bn`, `vgg19`, `vgg19_bn`, `cornet-s`, `clip-ViT`
@@ -64,7 +64,7 @@ You can find the jupyter notebook using `PyTorch` [here](https://colab.research.
 
 6. If you happen to extract hidden unit activations for many images, it is possible to run into `MemoryErrors`. To circumvent such problems, a helper function called `split_activations` will split the activation matrix into several batches, and stores them in separate files. For now, the split parameter is set to `10`. Hence, the function will split the activation matrix into `10` files. This parameter can, however, easily be modified in case you need more (or fewer) splits. To merge the separate activation batches back into a single activation matrix, just call `merge_activations` when loading the activations (e.g., `activations = merge_activations(PATH)`). 
 
-## Extract features at specific layer of a state-of-the-art `torchvision`, `TensorFlow`, `CORnet`, or `CLIP` model 
+## Extract features at specific layer of a state-of-the-art `torchvision`, `TensorFlow`, `CORnet`, or `CLIP`, `Timm` model 
 
 The following examples demonstrate how to load a model with PyTorch or TensorFlow into memory, and how to subsequently extract features. 
 Please keep in mind, that the model names as well as the layer names depend on the backend. If you use PyTorch, you will need to use these [model names](https://pytorch.org/vision/stable/models.html). If you use Tensorflow, you will need to use these [model names](https://keras.io/api/applications/). You can find the layer names by using `model.show()`. 
@@ -79,11 +79,11 @@ import thingsvision.vision as vision
 from thingsvision.model_class import Model
 
 model_name = 'alexnet'
-backend = 'pt'
+source = 'torchvision'
 batch_size = 64
 
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
-model = Model(model_name, pretrained=True, model_path=None, device=device, backend=backend)
+model = Model(model_name, pretrained=True, model_path=None, device=device, source=source)
 module_name = model.show()
 
 AlexNet(
@@ -123,7 +123,7 @@ dl = vision.load_dl(
 	out_path=f'./{model_name}/{module_name}/features',
 	batch_size=batch_size,
 	transforms=model.get_transformations(),
-	backend=backend,
+	backend=model.backend,
 	)
 features, targets, probas = model.extract_features(
 				data_loader=dl,
@@ -146,18 +146,17 @@ from thingsvision.model_class import Model
 
 model_name = 'clip-ViT'
 module_name = 'visual'
-backend = 'pt'
 batch_size = 64
 
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
-model = Model(model_name, pretrained=True, model_path=None, device=device, backend=backend)
+model = Model(model_name, pretrained=True, model_path=None, device=device)
 dl = vision.load_dl(
 	root='./images/',
 	out_path=f'./{model_name}/{module_name}/features',
 	batch_size=batch_size,
 	transforms=model.get_transformations(),
-	backend=backend,
+	backend=model.backend,
 	)
 features, targets = model.extract_features(
 			data_loader=dl,
@@ -180,11 +179,10 @@ import thingsvision.vision as vision
 from thingsvision.model_class import Model
 
 model_name = 'vit_b_16'
-backend = 'pt'
 batch_size = 64
 
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
-model = Model(model_name, pretrained=True, model_path=None, device=device, backend=backend)
+model = Model(model_name, pretrained=True, model_path=None, device=device)
 module_name = model.show()
 
 VisionTransformer(
@@ -389,7 +387,7 @@ dl = vision.load_dl(
 	out_path=f'./{model_name}/{module_name}/features',
 	batch_size=batch_size,
 	transforms=model.get_transformations(),
-	backend=backend,
+	backend=model.backend,
 	)
 
 features, targets, probas = model.extract_features(
@@ -507,7 +505,7 @@ dl = vision.load_dl(
 	out_path=f'./{model_name}/{module_name}/features',
 	batch_size=batch_size,
 	transforms=model.get_transformations(),
-	backend=backend,
+	backend=model.backend,
 	)
 features, targets = model.extract_features(
 			data_loader=dl,
@@ -530,19 +528,18 @@ import thingsvision.vision as vision
 from thingsvision.model_class import Model
 
 model_name = 'VGG16'
-backend = 'tf'
 module_name = 'block1_conv1'
 batch_size = 64
 
 device = 'cuda' if tf.test.is_gpu_available() else 'cpu'
-model = Model(model_name, pretrained=True, model_path=None, device=device, backend=backend)
+model = Model(model_name, pretrained=True, model_path=None, device=device)
 
 dl = vision.load_dl(
 	root='./images/',
 	out_path=f'./{model_name}/{module_name}/features',
 	batch_size=batch_size,
 	transforms=model.get_transformations(),
-	backend=backend,
+	backend=model.backend,
 	)
 features, targets, probas = model.extract_features(
 				data_loader=dl,
@@ -566,18 +563,18 @@ dl = vision.load_dl(
 		out_path=f'./{model_name}/{module_name}/features',
 		batch_size=batch_size,
 		transforms=model.get_transformations(apply_center_crop=apply_center_crop),
-		backend=backend,
+		backend=model.backend,
 		)
 ```
 
 ## Extract features from custom models
 
-If you want to use a custom model from the `custom_models` directory, you need to use the class name e.g. `VGG16_ecoset` as model name. The script will use the PyTorch or Tensorflow implementation depending on the `backend` value.
+If you want to use a custom model from the `custom_models` directory, you need to use the class name e.g. `VGG16_ecoset` as model name. 
 
 ```python
 from thingsvision.model_class import Model
 model_name = 'VGG16_ecoset'
-model = Model(model_name, pretrained=True, model_path=None, device=device, backend=backend)
+model = Model(model_name, pretrained=True, model_path=None, device=device)
 ```
 
 ## Representational Similarity Analysis (RSA) 
@@ -622,7 +619,6 @@ correlations = vision.compare_models(
                                      module_names=module_names,
                                      pretrained=True,
                                      batch_size=batch_size,
-                                     backend='pt',
                                      flatten_acts=True,
                                      clip=clip_list,
                                      save_features=True,
@@ -644,7 +640,7 @@ CLIP (Contrastive Language-Image Pre-Training) is a neural network trained on a
 ## Adding custom models
 
 If you want to use your own model and/or want to make it public, you just need to implement a class inheriting from the `custom_models/custom.py:Custom` class and implement the `create_model` method.
-There you can build/download the model and its weights. The constructors expects a `device` (str) and a `backend` (str).
+There you can build/download the model and its weights. The constructors expects a `device` (str).
 Afterwards you can put the file in the `custom_models` directory and create a pull request to include the model in the official GitHub repository.
 
 ```python
@@ -653,16 +649,15 @@ import torchvision.models as torchvision_models
 import torch
 
 class VGG16_ecoset(Custom):
-    def __init__(self, device, backend) -> None:
-        super().__init__(device, backend)
+    def __init__(self, device) -> None:
+        super().__init__(device)
 
     def create_model(self):
-        if self.backend == 'pt':
-            model = torchvision_models.vgg16_bn(pretrained=False, num_classes=565)
-            path_to_weights = 'https://osf.io/fe7s5/download'
-            state_dict = torch.hub.load_state_dict_from_url(path_to_weights, map_location=self.device)
-            model.load_state_dict(state_dict)
-            return model
+          model = torchvision_models.vgg16_bn(pretrained=False, num_classes=565)
+          path_to_weights = 'https://osf.io/fe7s5/download'
+          state_dict = torch.hub.load_state_dict_from_url(path_to_weights, map_location=self.device)
+          model.load_state_dict(state_dict)
+          return model
 ```
 
 ## Citation

doc/pytorch.ipynb

@@ -222,7 +222,6 @@
    },
    "outputs": [],
    "source": [
-    "backend = 'pt' # backend 'pt' for PyTorch or 'tf' for Tensorflow \n",
     "pretrained = True # use pretrained model weights\n",
     "model_path = None # if pretrained = False (i.e., randomly initialized weights) set path to model weights\n",
     "batch_size = 32 # use a power of two (this can be any size, depending on the number of images for which you aim to extract features)\n",
@@ -253,12 +252,15 @@
    "source": [
     "## load model\n",
     "model_name = 'vgg16_bn' \n",
+    "# specify the model source as VGG16 can be used from torchvision, timm, ... In this case torchvision is used (https://pytorch.org/vision/stable/models.html)\n",
+    "source = 'torchvision' \n",
+    "\n",
     "model = Model(\n",
     "            model_name,\n",
     "            pretrained=pretrained,\n",
     "            model_path=model_path,\n",
     "            device=device,\n",
-    "            backend=backend,\n",
+    "            source=source\n",
     ")"
    ]
   },
@@ -376,8 +378,7 @@
     "            model_name,\n",
     "            pretrained=pretrained,\n",
     "            model_path=model_path,\n",
-    "            device=device,\n",
-    "            backend=backend,\n",
+    "            device=device\n",
     ")"
    ]
   },
@@ -491,8 +492,7 @@
     "            model_name,\n",
     "            pretrained=pretrained,\n",
     "            model_path=model_path,\n",
-    "            device=device,\n",
-    "            backend=backend,\n",
+    "            device=device\n",
     ")"
    ]
   },

doc/tensorflow.ipynb

@@ -229,7 +229,6 @@
    },
    "outputs": [],
    "source": [
-    "backend = 'tf' # backend 'pt' for PyTorch or 'tf' for Tensorflow \n",
     "pretrained = True # use pretrained model weights\n",
     "model_path = None # if pretrained = False (i.e., randomly initialized weights) set path to model weights\n",
     "batch_size = 32 # use a power of two (this can be any size, depending on the number of images for which you aim to extract features)\n",
@@ -262,12 +261,15 @@
    "source": [
     "## load model\n",
     "model_name = 'VGG16' \n",
+    "# specify the model source, in this case use Keras applications (https://keras.io/api/applications/)\n",
+    "source = 'keras' \n",
+    "\n",
     "model = Model(\n",
     "            model_name,\n",
     "            pretrained=pretrained,\n",
     "            model_path=model_path,\n",
     "            device=device,\n",
-    "            backend=backend,\n",
+    "            source=source\n",
     ")"
    ]
   },

tests/helper.py

@@ -23,9 +23,9 @@
 TEST_PATH = './test_images'
 OUT_PATH = './test'
 
-PT_MODEL_AND_MODULE_NAMES = {
+MODEL_AND_MODULE_NAMES = {
     # Torchvision models
-    'vgg16_bn': {
+    'vgg16': {
         'modules': ['features.23', 'classifier.3'],
         'pretrained': True
     },
@@ -74,10 +74,8 @@
     'gluon_inception_v3': {
         'modules': ['Mixed_6d'],
         'pretrained': False
-    }
-}
+    },
 
-TF_MODEL_AND_MODULES_NAMES = {
     # Keras models
     'VGG16': {
         'modules': ['block1_conv1', 'flatten'],
@@ -89,7 +87,6 @@
     }
 }
 
-BACKENDS = ['tf', 'pt']
 
 FILE_FORMATS = ['hdf5', 'npy', 'mat', 'txt']
 DISTANCES = ['correlation', 'cosine', 'euclidean', 'gaussian']
@@ -160,34 +157,26 @@ def __len__(self) -> int:
         return len(self.values)
 
 def iterate_through_all_model_combinations():
-    for backend in BACKENDS:
-        MODEL_AND_MODULE_NAMES = None
-        if backend == 'pt':
-            MODEL_AND_MODULE_NAMES = PT_MODEL_AND_MODULE_NAMES
-        elif backend == 'tf':
-            MODEL_AND_MODULE_NAMES = TF_MODEL_AND_MODULES_NAMES
-
-        for model_name in MODEL_AND_MODULE_NAMES:
-            pretrained = MODEL_AND_MODULE_NAMES[model_name]['pretrained']
-            model, dataset, dl = create_model_and_dl(model_name, backend, pretrained)
+    for model_name in MODEL_AND_MODULE_NAMES:
+        pretrained = MODEL_AND_MODULE_NAMES[model_name]['pretrained']
+        model, dataset, dl = create_model_and_dl(model_name, pretrained)
 
-            modules =  MODEL_AND_MODULE_NAMES[model_name]['modules']
-            yield model, dataset, dl, modules, model_name
+        modules =  MODEL_AND_MODULE_NAMES[model_name]['modules']
+        yield model, dataset, dl, modules, model_name
 
 
-def create_model_and_dl(model_name, backend, pretrained=False):
-    """Iterate through all backends and models and create model, dataset and data loader."""
+def create_model_and_dl(model_name, pretrained=False):
+    """Iterate through models and create model, dataset and data loader."""
     model = Model(
         model_name=model_name,
         pretrained=pretrained,
-        device=DEVICE,
-        backend=backend
+        device=DEVICE
     )
 
     dataset = ImageDataset(
         root=TEST_PATH,
         out_path=OUT_PATH,
-        backend=backend,
+        backend=model.backend,
         imagenet_train=None,
         imagenet_val=None,
         things=None,
@@ -198,7 +187,7 @@ def create_model_and_dl(model_name, backend, pretrained=False):
     dl = DataLoader(
         dataset,
         batch_size=BATCH_SIZE,
-        backend=backend,
+        backend=model.backend,
     )
     return model, dataset, dl
 

tests/model/extraction/test_custom_model.py

@@ -18,16 +18,18 @@ def test_extract_features(self):
         backends = [['pt', helper.pt_model, 'vgg16'], ['tf', helper.tf_model, 'VGG16']]
         batch_size = 1
         for backend, custom_model, vgg_model in backends:
+            model = Model(vgg_model, pretrained=False,
+                        device=helper.DEVICE)
+            model.backend = backend
+            model.model = custom_model
+
             dataset = helper.SimpleDataset(values, backend)
             dl = DataLoader(
                 dataset,
                 batch_size=batch_size,
                 backend=backend,
             )
-            model = Model(vgg_model, pretrained=False,
-                        device=helper.DEVICE, backend=backend)
-
-            model.model = custom_model
+
             expected_features = np.array([[2, 2], [0, 0]])
             expected_targets = np.array([0, 0])
 
@@ -54,7 +56,8 @@ def test_extraction_batches(self):
         values = [1] * 10
         backend = 'pt'
         dataset = helper.SimpleDataset(values, backend)
-        model = Model('vgg16', pretrained=False, device=helper.DEVICE, backend=backend)
+        model = Model('vgg16', pretrained=False, device=helper.DEVICE)
+        model.backend = backend
         model.model = helper.pt_model
 
         # no batch remainders -> 5 batches with 2 examples