diff --git a/brainscore_vision/models/ReAlnet01/__init__.py b/brainscore_vision/models/ReAlnet01/__init__.py new file mode 100644 index 000000000..fab04f508 --- /dev/null +++ b/brainscore_vision/models/ReAlnet01/__init__.py @@ -0,0 +1,8 @@ +from brainscore_vision import model_registry +from brainscore_vision.model_helpers.brain_transformation import ModelCommitment +from .model import get_model, LAYERS + +model_registry['ReAlnet01']= lambda: ModelCommitment( + identifier='ReAlnet01', + activations_model=get_model(), + layers=LAYERS) diff --git a/brainscore_vision/models/ReAlnet01/model.py b/brainscore_vision/models/ReAlnet01/model.py new file mode 100644 index 000000000..4c5d6eb46 --- /dev/null +++ b/brainscore_vision/models/ReAlnet01/model.py @@ -0,0 +1,225 @@ +import math +from collections import OrderedDict +import torch +from torch import nn +from torchvision import transforms +import torch.utils.model_zoo +import os +from torch.utils.data import Dataset, DataLoader +import pandas as pd +import numpy as np +import torch.nn.functional as F +import h5py +import random + +import functools + +import torchvision.models + +import gdown + +from brainscore_vision.model_helpers.activations.pytorch import PytorchWrapper +from brainscore_vision.model_helpers.activations.pytorch import load_preprocess_images + + + +LAYERS = ['module.V1', 'module.V2', 'module.V4', 'module.IT', 'module.decoder.avgpool'] + +class Flatten(nn.Module): + + """ + Helper module for flattening input tensor to 1-D for the use in Linear modules + """ + + def forward(self, x): + return x.view(x.size(0), -1) + + +class Identity(nn.Module): + + """ + Helper module that stores the current tensor. Useful for accessing by name + """ + + def forward(self, x): + return x + + +class CORblock_S(nn.Module): + + scale = 4 # scale of the bottleneck convolution channels + + def __init__(self, in_channels, out_channels, times=1): + super().__init__() + + self.times = times + + self.conv_input = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False) + self.skip = nn.Conv2d(out_channels, out_channels, + kernel_size=1, stride=2, bias=False) + self.norm_skip = nn.BatchNorm2d(out_channels) + + self.conv1 = nn.Conv2d(out_channels, out_channels * self.scale, + kernel_size=1, bias=False) + self.nonlin1 = nn.ReLU(inplace=True) + + self.conv2 = nn.Conv2d(out_channels * self.scale, out_channels * self.scale, + kernel_size=3, stride=2, padding=1, bias=False) + self.nonlin2 = nn.ReLU(inplace=True) + + self.conv3 = nn.Conv2d(out_channels * self.scale, out_channels, + kernel_size=1, bias=False) + self.nonlin3 = nn.ReLU(inplace=True) + + self.output = Identity() # for an easy access to this block's output + + # need BatchNorm for each time step for training to work well + for t in range(self.times): + setattr(self, f'norm1_{t}', nn.BatchNorm2d(out_channels * self.scale)) + setattr(self, f'norm2_{t}', nn.BatchNorm2d(out_channels * self.scale)) + setattr(self, f'norm3_{t}', nn.BatchNorm2d(out_channels)) + + def forward(self, inp): + x = self.conv_input(inp) + + for t in range(self.times): + if t == 0: + skip = self.norm_skip(self.skip(x)) + self.conv2.stride = (2, 2) + else: + skip = x + self.conv2.stride = (1, 1) + + x = self.conv1(x) + x = getattr(self, f'norm1_{t}')(x) + x = self.nonlin1(x) + + x = self.conv2(x) + x = getattr(self, f'norm2_{t}')(x) + x = self.nonlin2(x) + + x = self.conv3(x) + x = getattr(self, f'norm3_{t}')(x) + + x += skip + x = self.nonlin3(x) + output = self.output(x) + + return output + + +def CORnet_S(): + model = nn.Sequential(OrderedDict([ + ('V1', nn.Sequential(OrderedDict([ # this one is custom to save GPU memory + ('conv1', nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, + bias=False)), + ('norm1', nn.BatchNorm2d(64)), + ('nonlin1', nn.ReLU(inplace=True)), + ('pool', nn.MaxPool2d(kernel_size=3, stride=2, padding=1)), + ('conv2', nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1, + bias=False)), + ('norm2', nn.BatchNorm2d(64)), + ('nonlin2', nn.ReLU(inplace=True)), + ('output', Identity()) + ]))), + ('V2', CORblock_S(64, 128, times=2)), + ('V4', CORblock_S(128, 256, times=4)), + ('IT', CORblock_S(256, 512, times=2)), + ('decoder', nn.Sequential(OrderedDict([ + ('avgpool', nn.AdaptiveAvgPool2d(1)), + ('flatten', Flatten()), + ('linear', nn.Linear(512, 1000)), + ('output', Identity()) + ]))) + ])) + + # weight initialization + for m in model.modules(): + if isinstance(m, nn.Conv2d): + n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + m.weight.data.normal_(0, math.sqrt(2. / n)) + # nn.Linear is missing here because I originally forgot + # to add it during the training of this network + elif isinstance(m, nn.BatchNorm2d): + m.weight.data.fill_(1) + m.bias.data.zero_() + + return model + + +class Encoder(nn.Module): + def __init__(self, realnet, n_output): + super(Encoder, self).__init__() + + # CORnet + self.realnet = realnet + + # full connected layer + self.fc_v1 = nn.Linear(200704, 128) + self.fc_v2 = nn.Linear(100352, 128) + self.fc_v4 = nn.Linear(50176, 128) + self.fc_it = nn.Linear(25088, 128) + self.fc = nn.Linear(512, n_output) + self.activation = nn.ReLU() + + def forward(self, imgs): + + outputs = self.realnet(imgs) + + N = len(imgs) + v1_outputs = self.realnet.module.V1(imgs) # N * 64 * 56 * 56 + v2_outputs = self.realnet.module.V2(v1_outputs) # N * 128 * 28 * 28 + v4_outputs = self.realnet.module.V4(v2_outputs) # N * 256 * 14 * 14 + it_outputs = self.realnet.module.IT(v4_outputs) # N * 512 * 7 * 7 + v1_features = self.fc_v1(v1_outputs.view(N, -1)) + v1_features = self.activation(v1_features) + v2_features = self.fc_v2(v2_outputs.view(N, -1)) + v2_features = self.activation(v2_features) + v4_features = self.fc_v4(v4_outputs.view(N, -1)) + v4_features = self.activation(v4_features) + it_features = self.fc_it(it_outputs.view(N, -1)) + it_features = self.activation(it_features) + features = torch.cat((v1_features, v2_features, v4_features, it_features), dim=1) + features = self.fc(features) + + return outputs, features + + +torch.set_default_dtype(torch.float32) + +transform = transforms.Compose( + [ + transforms.Resize((224, 224)), + transforms.ToTensor(), + transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) + ] +) + +device='cuda' + +realnet = CORnet_S().to(device) +realnet = torch.nn.DataParallel(realnet) +encoder = Encoder(realnet, 340).to(device) +# weights = torch.load('/Users/yat-lok/Desktop/ReAInet/ReAInet_weights/sub-01/best_model_params.pt') +# weights = torch.load('/home/yilwang_umass_edu/ReAlnet_weights/sub-01/best_model_params.pt', map_location=device) +# weights = torch.load('/work/pi_gstuart_umass_edu/yile/weights_1/sub-01.pt', map_location=device) +# weights = torch.load('/work/pi_gstuart_umass_edu/yile/weights_10/sub-01.pt', map_location=device) + +url = 'https://drive.google.com/uc?id=1AtpS7dPV8t3e1aT8a4Nu-mIFkbWRH-ff' +output_file = "best_model_params.pt" +gdown.download(url, output_file, quiet = False) + +weights = torch.load(output_file, map_location=device) + + +encoder.load_state_dict(weights) +realnet = encoder.realnet + +realnet.eval() + +def get_model(): + model = realnet + preprocessing = functools.partial(load_preprocess_images, image_size=224) + wrapper = PytorchWrapper(identifier='ReAlnet01', model=model, preprocessing=preprocessing) + wrapper.image_size = 224 + return wrapper diff --git a/brainscore_vision/models/ReAlnet01/requirments.txt b/brainscore_vision/models/ReAlnet01/requirments.txt new file mode 100644 index 000000000..3b5cb3446 --- /dev/null +++ b/brainscore_vision/models/ReAlnet01/requirments.txt @@ -0,0 +1,7 @@ + +torch +torchvision +pandas +numpy +h5py +gdown diff --git a/brainscore_vision/models/ReAlnet01/setup.py b/brainscore_vision/models/ReAlnet01/setup.py new file mode 100644 index 000000000..fe427ad08 --- /dev/null +++ b/brainscore_vision/models/ReAlnet01/setup.py @@ -0,0 +1,18 @@ +from setuptools import setup + +setup( + name='ReAlnet01', + version='1.0.0', + author='Zitong Lu', + install_requires=[ + 'math', + 'collections', + 'torch', + 'torchvision', + 'os', + 'pandas', + 'numpy', + 'h5py', + 'random', + 'gdown'], +) diff --git a/brainscore_vision/models/ReAlnet01/test.py b/brainscore_vision/models/ReAlnet01/test.py new file mode 100644 index 000000000..badc557ae --- /dev/null +++ b/brainscore_vision/models/ReAlnet01/test.py @@ -0,0 +1,15 @@ +import pytest +from pytest import approx + +from brainscore_vision import score + + +@pytest.mark.private_access +@pytest.mark.memory_intense +@pytest.mark.parametrize("model_identifier, benchmark_identifier, expected_score", [ + ("ReAlnet01", "MajajHong2015.IT-pls", approx(0.0153, abs=0.0005)), +]) +def test_score(model_identifier, benchmark_identifier, expected_score): + actual_score = score(model_identifier=model_identifier, benchmark_identifier=benchmark_identifier, + conda_active=True) + assert actual_score == expected_score