Skip to content

Commit

Permalink
Merge branch 'master' into keepDevelopBranchSynchronized
Browse files Browse the repository at this point in the history
  • Loading branch information
KartikP authored Aug 26, 2024
2 parents f0b94b5 + ed16fd9 commit 25c540a
Show file tree
Hide file tree
Showing 4 changed files with 119 additions and 0 deletions.
9 changes: 9 additions & 0 deletions brainscore_vision/models/eBarlow_lmda_02_1000ep/__init__.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,9 @@
from brainscore_vision import model_registry
from brainscore_vision.model_helpers.brain_transformation import ModelCommitment
from .model import get_model, get_layers

model_registry["eBarlow_lmda_02_1000ep"] = lambda: ModelCommitment(
identifier="eBarlow_lmda_02_1000ep",
activations_model=get_model("eBarlow_lmda_02_1000ep"),
layers=get_layers("eBarlow_lmda_02_1000ep"),
)
84 changes: 84 additions & 0 deletions brainscore_vision/models/eBarlow_lmda_02_1000ep/model.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,84 @@
from brainscore_vision.model_helpers.check_submission import check_models
import functools
import os
from urllib.request import urlretrieve
import torchvision.models
from brainscore_vision.model_helpers.activations.pytorch import PytorchWrapper
from brainscore_vision.model_helpers.activations.pytorch import load_preprocess_images
from pathlib import Path
from brainscore_vision.model_helpers import download_weights
import torch
from collections import OrderedDict

# This is an example implementation for submitting resnet-50 as a pytorch model

# Attention: It is important, that the wrapper identifier is unique per model!
# The results will otherwise be the same due to brain-scores internal result caching mechanism.
# Please load your pytorch model for usage in CPU. There won't be GPUs available for scoring your model.
# If the model requires a GPU, contact the brain-score team directly.
from brainscore_vision.model_helpers.check_submission import check_models


def get_model_list():
return ["eBarlow_lmda_02_1000ep"]


def get_model(name):
assert name == "eBarlow_lmda_02_1000ep"
url = "https://users.flatironinstitute.org/~tyerxa/equi_proj/training_checkpoints/fresh/paired/lmda_0.2/Barlow_1000ep/ep1000-ba625000-rank0"
fh = urlretrieve(url)
state_dict = torch.load(fh[0], map_location=torch.device("cpu"))["state"]["model"]
model = load_composer_classifier(state_dict)
preprocessing = functools.partial(load_preprocess_images, image_size=224)
wrapper = PytorchWrapper(identifier=name, model=model, preprocessing=preprocessing)
wrapper.image_size = 224
return wrapper

def load_composer_classifier(sd):
model = torchvision.models.resnet.resnet50()
new_sd = OrderedDict()
for k, v in sd.items():
if 'lin_cls' in k:
new_sd['fc.' + k.split('.')[-1]] = v
if ".f." not in k:
continue
parts = k.split(".")
idx = parts.index("f")
new_k = ".".join(parts[idx + 1 :])
new_sd[new_k] = v
model.load_state_dict(new_sd, strict=True)
return model

def get_layers(name):
assert name == "eBarlow_lmda_02_1000ep"
layers = [
"layer1.0",
"layer1.1",
"layer1.2",
"layer2.0",
"layer2.1",
"layer2.2",
"layer2.3",
"layer3.0",
"layer3.1",
"layer3.2",
"layer3.3",
"layer3.4",
"layer3.5",
"layer4.0",
"layer4.1",
"layer4.2",
"avgpool",
"fc",
]
outs = ["conv1", "layer1", "layer2", "layer3", "layer4"]

return layers + outs


def get_bibtex(model_identifier):
return """xx"""


if __name__ == "__main__":
check_models.check_base_models(__name__)
25 changes: 25 additions & 0 deletions brainscore_vision/models/eBarlow_lmda_02_1000ep/setup.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,25 @@
#!/usr/bin/env python
# -*- coding: utf-8 -*-

from setuptools import setup, find_packages

requirements = [ "torchvision",
"torch"
]

setup(
packages=find_packages(exclude=['tests']),
include_package_data=True,
install_requires=requirements,
license="MIT license",
zip_safe=False,
keywords='brain-score template',
classifiers=[
'Development Status :: 2 - Pre-Alpha',
'Intended Audience :: Developers',
'License :: OSI Approved :: MIT License',
'Natural Language :: English',
'Programming Language :: Python :: 3.7',
],
test_suite='tests',
)
1 change: 1 addition & 0 deletions brainscore_vision/models/eBarlow_lmda_02_1000ep/test.py
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
# Left empty as part of 2023 models migration

0 comments on commit 25c540a

Please sign in to comment.