Custom Models
Deep Dive 3 will go over how to create and submit your own custom model package. This Deep Dive assumes that you have completed both the Quickstart Tutorial, Deep Dive 1, and Deep Dive 2 so if you have not yet completed those, please do so before continuing.
Part 1 will discuss how the tutorial model package from Deep Dive 1 can be modified to create your own custom model. In particular, we'll look at a single required change in the __init__.py file, and more extensive changes required in the model.py file. Part 2 will cover some resources that might be beneficial in creating your own models, and finally, Part 3 will provide directions on how to submit your model through the BrainScore website, as well as directly through the BrainScore GitHub repository.
NOTE: Please do not actually submit the model that you prepare in this tutorial.
Part 1: Defining the Model Itself
The submission package for a custom model is nearly identical to the sample resnet50_tutorial submission folder we examined in detail in Deep Dive #1. The only changes will be to the model.py and the __init__.py files as we need to first define our model and then add it to the Brain-Score ecosystem.
First, regarding __init__.py file; this file is discussed in detail in Deep Dive 1, and as stated there, it's important to ensure that the identifier you specify in for the model_registry is unique.
Next, let's look at the model.py file. The model.py file in the sample resnet50_tutorial submission folder that we looked at in Deep Dive 1 was used to load a pre-trained model. Instead of loading a pre-trained model, we show here how to use the model.py file to define a completely new model.
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. Next, lines 18-33 define the custom neural network model called MyCustomModel using PyTorch.
The init function on lines 19-25 initializes the model by setting up its layers. More specifically, in line 21, a convolution layer is created which takes a stimuli with 3 channels (typically RGB), applies 2 filters, and uses a 3x3 kernel. Line 22 uses a ReLU activation function to introduce non-linearity after convolution. Line 23 calculates the input size for the fully connected layer based on the output of the convolution layer (assuming an input image size of 224x2). In line 24, a fully connected layer connects the flattened output of the convolution layer to a 1000-dimensional output. How the conv1 layer is producing a flattened output is described in Line 30 in the forward function using x.view(x.size(0)). Finally, in line 25, another ReLU activation function is used to transform the output of the fully connected layer.
Line 42 is a call to Brain-Score's built in vision model preprocessing. 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.
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.
Part 2: Additional Resources for Creating Your Model
When creating a custom model, we recommend following the guidelines and processes for creating PyTorch models. An excellent tutorial can be found on the PyTorch model creation page.
This page gives background on the types of model layers available (such as linear, convolutional, recurrent, or data manipulation layers) as well as different types of activation functions, and even some discussion of advanced architectures like transformers.
NOTE: If you adhere to PyTorch's recommended procedures, your model should be compatible with BrainScore. However, please ensure that the outputs of any layers for which you want BrainScore to evaluate layer activations are formatted as tensors. It's important to note that while PyTorch models can be created such that layers output data in other formats, such as tuples or numpy arrays, BrainScore requires tensor-formatted layer outputs for proper evaluation.
Part 3: How To Submit Your Completed Model
There are two ways to submit your completed model. The first is through the BrainScore website and the second is directly through the BrainScore Github repository for the appropriate domain (Vision or Language). If you submit through the website, your submission package must be set up with the proper directory structure described in Part 3 of Deep Dive #1, and then compressed into a zip file. If you submit through a Github PR, then you just need the model folder containing the four files (also) described in Deep Dive #1.
To submit directly through Github, prepare your submission folder with the four files discussed in Deep Dive #1 Then place this folder in /models (for Vision) or /models (for Language). Then simply open a Pull Request and a member of the Brain-Score team will review it and either merge it or suggest changes. Once this PR is merged, it will be scored, and you will receive an email containing those scores. NOTE: If you are submitting a plugin without other changes and don't need a review, you can use the automerge tag with your PR and it will be merged automatically after merging tests.
The second way to submit is through the BrainScore website. Click on either Vision or Language, and then click on the "Submit" button on the right hand side of the page, and follow the instructions to submit. When submitting through the website, make sure that your zip file contains the folder structure showin in Part 3 of Deep Dive 2. NOTE: You will have to first create a BrainScore account and be logged in to submit a custom model.
Our tutorials and FAQs, created with Brain-Score users, aim to cover all bases. However, if issues arise, reach out to our community or consult the troubleshooting guide below for common errors and solutions.