Deep Dive 2

Part 1: Defining the Model Itself

Luckily, the submission package for a custom model is nearly identical to our sample project Zip that we explored in Deep Dive 1. The only changes will be to the __init__.py file model.py files, as we need to define our model and add it into the Brain-Score ecosystem.

We will start our exploration of a custom model submission via the model.py file, which actually contains the model. Let's take a look below at this file...

1       from brainscore_vision.model_helpers.check_submission import check_models
2       import functools
3       from brainscore_vision.model_helpers.activations.pytorch import PytorchWrapper
4       from brainscore_vision.model_helpers.activations.pytorch import load_preprocess_images
5       import torch
6       import numpy as np
7       from brainscore_vision.model_helpers.brain_transformation import ModelCommitment
8
9
10      # This is an example implementation for submitting a custom model named my_custom_model
11
12      # Attention: It is important that the wrapper identifier is unique per model!
13      # You will be unable to submit a model with the same identifier as an existing model.
14
15      # As a reminder, please do not actually submit this model to the Brain-Score platform
16
17
18      class MyCustomModel(torch.nn.Module):
19          def __init__(self):
20              super(MyCustomModel, self).__init__()
21              self.conv1 = torch.nn.Conv2d(in_channels=3, out_channels=2, kernel_size=3)
22              self.relu1 = torch.nn.ReLU()
23              linear_input_size = np.power((224 - 3 + 2 * 0) / 1 + 1, 2) * 2
24              self.linear = torch.nn.Linear(int(linear_input_size), 1000)
25              self.relu2 = torch.nn.ReLU()  # can't get named ReLU output otherwise
26
27          def forward(self, x):
28              x = self.conv1(x)
29              x = self.relu1(x)
30              x = x.view(x.size(0), -1)
31              x = self.linear(x)
32              x = self.relu2(x)
33              return x
34
35
36      def get_model_list():
37          return ['my_custom_model']
38
39
40      def get_model(name):
41          assert name == 'my_custom_model'
42          preprocessing = functools.partial(load_preprocess_images, image_size=224)
43          activations_model = PytorchWrapper(identifier='my_custom_model', model=MyCustomModel(), preprocessing=preprocessing)
44          model = ModelCommitment(identifier='my_custom_model', activations_model=activations_model,
45                                  # specify layers to consider
46                                  layers=['conv1', 'relu1', 'relu2'])
47          wrapper.image_size = 224
48          return wrapper
49
50
51      def get_layers(name):
52          assert name == 'my_custom_model'
53          return ['conv1', 'relu1', 'relu2']
54
55
56      def get_bibtex(model_identifier):
57          return """xx"""
58
59
60      if __name__ == '__main__':
61          check_models.check_base_models(__name__)
    

Lines 1-15 are the standard imports needed for the Brain-Score library, as well as some comments regarding model GPU uses. Next, lines 18-33 define the model itself, via a class named MyCustomModel. This is where you would add in your custom model, or you can simply import it if your model is complex/located in a different file.

Lines 36-37 tells Brain-Score which models to look at. You can add as many models as you would like, but just make sure to add their names here. As we only have one model, which we named my_custom_model, this is the only model defined in the list on line 37.

Line 42 is a call to Brain-Score's built in vision model preprocessing. This is standard, and you should not have to worry too much about this. Next up are line 43 and 44, which allow us to actually get the activations model from our PyTorch Wrapper and to make the model into a Brain-Score ModelCommitment, respectively. See the below section for why we need the ModelCommitment.

Lastly, lines 51-53 are where you tell Brain-Score which model layers to include, and lines 56-57 are where you would add in your model's Bibtex.