torch_multiscale_feature.py

# uvx --with torch --with torchvision ipython
from torchvision.models import get_model
model = get_model('resnet50', pretrained=True)

# https://pytorch.org/vision/main/feature_extraction.html
from torchvision.models.feature_extraction import create_feature_extractor
return_nodes = {f'layer{i}': f'layer{i}' for i in range(1, 5)}
extractor = create_feature_extractor(model, return_nodes)

import torch
img = torch.randn((3, 3, 256, 256))

features = extractor(img)
for i in range(1, 5):
    print(i, features[f'layer{i}'].shape)