Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Allow region_layer_map json to be loaded into BrainModel during commitment #1314

Merged
merged 8 commits into from
Oct 28, 2024
Merged

Allow region_layer_map json to be loaded into BrainModel during commitment #1314

Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
af8b25f
Add load_region_layer_map_from_json
KartikP Oct 10, 2024
938b460
Cleaned up load_region_layer_map_json function and added error handling.
KartikP Oct 10, 2024
b7fce8d
Merge branch 'master' into kp/model-commitment-map-retrieval
KartikP Oct 10, 2024
bc24c90
Replaced print with logging
KartikP Oct 12, 2024
4acf314
Recommended fix:
KartikP Oct 17, 2024
579e888
Recommended fix:
KartikP Oct 17, 2024
a43b281
Merge branch 'master' into kp/model-commitment-map-retrieval
KartikP Oct 17, 2024
224ea92
Merge branch 'master' into kp/model-commitment-map-retrieval
KartikP Oct 28, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 33 additions & 0 deletions brainscore_vision/model_helpers/brain_transformation/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,3 +1,8 @@
import json
from pathlib import Path
import logging
from brainscore_vision.utils import fullname
from brainscore_core.plugin_management import import_plugin
from brainscore_vision import load_benchmark
from brainscore_vision.model_helpers.brain_transformation.temporal import TemporalAligned
from brainscore_vision.model_interface import BrainModel
Expand All @@ -22,6 +27,7 @@ class ModelCommitment(BrainModel):
def __init__(self, identifier,
activations_model, layers, behavioral_readout_layer=None, region_layer_map=None,
visual_degrees=8):
self._logger = logging.getLogger(fullname(self))
self.layers = layers
self.activations_model = activations_model
# We set the visual degrees of the ActivationsExtractorHelper here to avoid changing its signature.
Expand All @@ -30,12 +36,18 @@ def __init__(self, identifier,
self.activations_model._extractor.set_visual_degrees(visual_degrees) # for microsaccades
self._visual_degrees = visual_degrees
# region-layer mapping

# Attempt to load region_layer_map from JSON, if available
region_layer_map = self.load_region_layer_map_json(identifier) if region_layer_map is None else region_layer_map
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't like that the ModelCommitment has to know about plugins. Can we not load the specification inside the plugin and pass it to the ModelCommitment as a parameter?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Modifying the __init__.py was the original approach. However, after talking with the others, it seems modifying the init file, especially when there can be multiple model identifiers listed in a single model directory could be messy. It would also make for a messier PR/commit during the model submission.

By adding a check in ModelCommitment, all that would be needed would be a single region_layer_map json per model_identifier in the model directory. But I do agree that it disrupts the current flow of the plugin model.

If we do not want to modify the __init__.py file of each model (and therefore cannot pass the map to ModelCommitment as a parameter), what would you recommend?

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

just thinking out loud, but what about an intermediate approach where the ModelCommitment expects the full layer mapping, but submissions can specify this to be automatically discovered? Something like

model_registry['mymodel'] = lambda: ModelCommitment(
    identifier='mymodel',
    activations_model=get_model(),
    region_layer_map=discover_region_layer_mapping('mymodel'))

The discover_region_layer_mapping would then first look for a mapping file in the plugin, and if that doesn't exist make one by running the layer selection. This way the plugin __init__ would not have to be modified but we still maintain the separation between plugins and model helpers.

Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Martin - Kartik and I just discussed and I would like to propose something a bit more direct that what you suggested. We create a means to determine if we are running locally (any place but AWS) or in production (AWS). For local runs (again, anything but AWS) if no json file exists, we perform layer mapping and then write the json file.
For AWS runs we assume that region_layer_map json file exists and throw an exception if it does not, terminating the scoring run.

This is changing the assumptions of modelCommitment, which hopefully we don't do frequently, but not significantly (and not even as significantly as the solution you propose above). Additionally, as a benefit, local users, at the time of their PR already have the layer map and have thought about it and checked it against reality. They can regenerate as many times as desired, simply by removing the json file.

Copy link
Member

@mschrimpf mschrimpf Oct 25, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

yes I think we agree on those points.

What I'm not sure about is if you want this detection to happen inside ModelCommitment or in the plugin. I am arguing for the latter so that we keep the clean decoupling between the plugin system and model helpers. The only advantage I see for putting it inside the ModelCommitment is to avoid the model plugins having to specify layer discovery, but I like that it makes this explicit to users. And the downside of having it inside of ModelCommitment is that it makes adapting this mapping more difficult for models such as CORnet and VOneNet.


# If region_layer_map is unavailable
if region_layer_map is None:
layer_selection = LayerSelection(model_identifier=identifier,
activations_model=activations_model, layers=layers,
visual_degrees=visual_degrees)
region_layer_map = RegionLayerMap(layer_selection=layer_selection,
region_benchmarks=STANDARD_REGION_BENCHMARKS)
Comment on lines 44 to 49
Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Write local JSON file to the model directory. This can be PR'd when they choose to submit.


# neural
layer_model = LayerMappedModel(identifier=identifier, activations_model=activations_model,
region_layer_map=region_layer_map)
Expand All @@ -52,6 +64,27 @@ def __init__(self, identifier,
})
self.do_behavior = False

def load_region_layer_map_json(self, identifier):
'''
Attempts to load the region_layer_map from a JSON file in the model's directory
If file exists, load JSON. Otherwise, return None and proceed with legacy layer mapping
'''
try:
importer = import_plugin.ImportPlugin(library_root='brainscore_vision', plugin_type='models', identifier=identifier)
model_dir = importer.locate_plugin()
project_root = Path(__file__).resolve().parent.parent
region_layer_map_path = project_root / 'vision' / 'models' / model_dir / 'region_layer_map' / f'{identifier}.json'
if region_layer_map_path.exists():
with region_layer_map_path.open('r') as region_layer_map_file:
self._logger.info(f"Successfully loaded region_layer_map for {identifier}")
return json.load(region_layer_map_file)
else:
self._logger.info(f"No region_layer_map file found for {identifier}, proceeding with default layer mapping")
return None
except Exception as e:
self._logger.error(f"Error importing model to search for region_layer_map: {e}")
return None

def visual_degrees(self) -> int:
return self._visual_degrees

Expand Down
Loading