brain-score / brain-score/vision
Is it possible to pass a model a torch.Tensor and return a torch.Tensor?
- 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