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

[billbrod]

I am trying to figure out how to use brainscore's models with 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:

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 to use with plenoptic, with an optional image transform and specific layers (using torchvision's feature extractor):

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

[mike-ferguson]

Hi @billbrod- I am looping in @mschrimpf here who I think will be able to give some more insight

[mschrimpf]

hi @billbrod, indeed the conversion to numpy in _tensor_to_numpy breaks the propagation of tensors:

vision/brainscore_vision/model_helpers/activations/pytorch.py

Lines 81 to 82 in 6c09bdf

	def _tensor_to_numpy(cls, output):
	return output.cpu().data.numpy()

I think the easiest would be to fork the repository and change this. Alternatively you could change this during runtime, e.g. something like

from brainscore_vision.model_helpers.activations.pytorch import PytorchWrapper

def tensor_stays_tensor(cls, output):
    return output

PytorchWrapper._tensor_to_numpy = tensor_stays_tensor

[billbrod]

Thanks @mschrimpf ! With the following (note that I removed the cls input from tensor_stays_tensor), I get a tensor output from the model.

from brainscore_vision.model_helpers.activations.pytorch import PytorchWrapper

def tensor_stays_tensor(output):
    return output

PytorchWrapper._tensor_to_numpy = tensor_stays_tensor

However, that still doesn't allow me to use the model with plenoptic, because there's a with torch.no_grad line in get_activations() as well (plenoptic needs to differentiate through the model).

But after looking at the source code for a little bit, I realized that what I actually want is model.activations_model._model. Doing the following allows me to use the model with plenoptic:

import plenoptic as po
import brainscore_vision
import torch
from torchvision.models.feature_extraction import create_feature_extractor

class TorchVision(torch.nn.Module):
    def __init__(self, model, return_node: str, transform = None):
        super().__init__()
        self.transform = transform
        self.extractor = create_feature_extractor(model, return_nodes=[return_node])
        self.model = model
        self.return_node = return_node
    
    def forward(self, x):
        if self.transform is not None:
            x = self.transform(x)
        return self.extractor(x)[self.return_node]

model = brainscore_vision.load_model("alexnet_training_seed_01")
# or some other layer
plenoptic_model = TorchVision(model.activations_model._model, model.layers[0])
po.tools.remove_grad(plenoptic_model)
img = po.data.einstein(as_gray=False).to(0)
plenoptic_model.to(0)
plenoptic_model.eval()
met = po.synth.Metamer(img, plenoptic_model)
met.synthesize(500)

It looks like the imagenet transform is included in model.activations_models._model, because img in the example above has shape (1,3,256,256) (and is un-normalized), so the cropping (and hopefully normalization) is being handled in the model. That's correct, right?

I just want to make sure -- if I use the model as above, I'm using the specific "BrainScore version" of the model, right? It looks to me like most of the stuff handled by PytorchWrapper is mainly handling multiple layers and getting things in the right shape (which I can handle with the torchvision extractor and in my little class above), is that correct?

[mschrimpf]

Indeed, model.activations_model._model will give you the raw pytorch model.

But note that the preprocessing is handled in the ActivationsExtractor class:

vision/brainscore_vision/model_helpers/activations/core.py

Line 200 in 6c09bdf

preprocessed_inputs = self.preprocess(inputs)

not the raw pytorch _model. You could access this via model.activations_model._extractor.preprocess

Also note that a small number of models (e.g. stochastic variants of VOneNet) have noise in their processing, so the multiple trials in these models actually matter:

vision/brainscore_vision/models/VOneCORnet_S/helpers/vonecornets.py

Line 18 in 6c09bdf

for trial_number in range(number_of_trials):

You can probably ignore these models for your project, just wanted to make sure you're aware.

[billbrod]

Thanks for your response. For the model I'm testing with, it looks like model.activations_model._extractor.preprocess is running load_preprocess_images, which loads the image from a filepath and then applies the preprocessing (resize and normalize). So I can't pass a tensor to this function -- it only works on filepaths. Looking through the source code, I figured the function I wanted was brainscore_vision.model_helpers.activations.pytorch.preprocess_images, which is called by load_preprocess_images after the load. However, this function also cannot accept a tensor, because it includes the ToTensor() transform as one of the steps. This transform can operate on either a numpy array or a PIL image, but you call it after the Resize transform, which can only operate on PIL images or torch tensors -- so only a PIL image will work with the preprocess_images function.

Given that the only relevant steps of this preprocessing for me are the resizing and the normalizing, I can just make sure to apply those steps myself. It looks like the there are two possible ways the preprocessing can go: either imagenet or inception, as specified by the preprocess_type arg. I'm assuming that depends on the model -- is that info stored somewhere so that I can grab it?

You can probably ignore these models for your project, just wanted to make sure you're aware.

Thanks for the heads up! Agreed it makes sense to ignore those models

[mschrimpf]

It looks like the there are two possible ways the preprocessing can go: either imagenet or inception, as specified by the preprocess_type arg. I'm assuming that depends on the model -- is that info stored somewhere so that I can grab it?

Yes, the models have to specify it: https://github.com/search?q=repo%3Abrain-score%2Fvision+preprocess_type+&type=code

[billbrod]

I see, so it's specified in the model.py for each model. But it doesn't look like that info is stored in any user-accessible attribute of the model, so I'll need to look at the source code for each model in order to determine the proper preprocessing? (Same with the expected image size, it looks like)

[mschrimpf]

Yea unfortunately. FWIW the image_size is often also stored in the wrapper:

vision/brainscore_vision/models/alexnet/model.py

Line 27 in 936d6c4

wrapper.image_size = 224

And from what I've seen preprocessing is usually just the same standard implementation. I guess you could run a dummy image through to verify.

[billbrod]

Gotcha.

As I played around with this, I think there's something I'm not understanding about how image size is handled though. With the alexnet model, which expects images with height and width 224, I can pass it images of any size:

>>> import torch
>>> import brainscore_vision
>>> model = brainscore_vision.load_model("alexnet")
>>> model.activations_model.image_size
224
>>> model.activations_model._model(torch.rand(1,3,1000,1000)).shape
torch.Size([1, 1000])
>>> model.activations_model._model(torch.rand(1,3,224,224)).shape
torch.Size([1, 1000])

So it looks like the cropping is being done in the forward pass -- or is there something else going on? Is the normalization? (I'm not sure how to check that as easily...)

[mschrimpf]

Not sure without looking in detail sorry -- if it was only convolutional layers, it should be size-agnostic but it looks like you're processing all the way through the fully-connected layers too