brain-score / brain-score/vision

Is it possible to pass a model a torch.Tensor and return a torch.Tensor?

Open
#1,848 10 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
193
Forks
105
Avg merge
10h 48m
Merged PRs (30d)
10

Description

I am trying to figure out how to use brainscore's models with [plenoptic](https://github.com/plenoptic-org/plenoptic/). Plenoptic requires a model that accepts a 4d torch.Tensor of images (shape `batch, channel, height, width`) and returns a 3d or 4d torch.Tensor of images (in a way that torch knows how to autodiff). I cannot figure out how to interact with brainscore's models in that way. Is there a way to do so?

Here is what I've tried:

```python
import brainscore_vision
model = brainscore_vision.load_model("alexnet_training_seed_01")
img = torch.rand(1,3,256,256)
rep = model.activations_model.get_activations(img, ["features.0"])["features.0"]
```

`rep` is then a numpy array (the shape is fine), not a torch tensor. Looking at `model_helpers/activations/pytorch.py`, it looks like this is because `_tensor_to_numpy` is registered as a hook -- is it possible to disable this hook?

As an example of what I'm talking about, here's how one would prepare a model from [timm](https://timm.fast.ai/) to use with plenoptic, with an optional image transform and specific layers (using torchvision's feature extractor):

```python
import timm
from timm.data import resolve_data_config
from timm.data.transforms_factory import create_transform
from torchvision.models.feature_extraction import create_feature_extractor

class IntermediateOutputResnet(nn.Module):
def __init__(self, model: nn.Module, return_node: str, transform: Optional[Callable] = None):
super().__init__()
self.return_node = return_node
self.extractor = create_feature_extractor(model, return_nodes=[return_node])
self.model = model
self.transform = transform

def forward(self, x):
if self.transform is not None:
x = self.transform(x)
return self.extractor(x)[self.return_node]

model = timm.create_model("hf-hub:nateraw/resnet50-oxford-iiit-pet", pretrained=True)
model.eval()
transform = create_transform(**resolve_data_config(model.pretrained_cfg, model=model))
test_model = IntermediateOutputResnet(model, "layer2", transform)

test_model(torch.rand(1,3,256,256))
# returns a 4d tensor
```

Contributor guide

No contributing guide indexed for this repository

Research direction

Start by reading model_helpers/activations/pytorch.py and tracing brainscore_vision.load_model through activations_model.get_activations. Compare the tensor-to-NumPy hook behavior with plenoptic's requirement for differentiable 4D tensor inputs and 3D or 4D tensor outputs. Done means the supported interaction and its scope are clearly established, with relevant behavior covered by tests.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.