thingsvision

thingsvision is a Python package designed for extracting image representations (features) from various state-of-the-art deep neural networks.

Core Workflow:

1. Provide a directory containing images.
2. Specify a neural network model and its source.
3. thingsvision processes the images and returns a feature vector for each image.

These feature vectors can then be used for downstream computer vision analyses.

For interactive exploration, you can use the provided Colab notebooks for PyTorch or Tensorflow.

You can extract features from a wide variety of model architectures and sources, including:

• Standard Libraries: torchvision, Keras, timm.
• Self-Supervised Learning (SSL): Models like simclr-rn50, mocov2-rn50, dino-vit-{tiny/small/base}-p{8/16}, dinov2-vit-{small/base/large/giant}-p14, and mae-vit variants.
• CLIP-based Models: OpenCLIP (LAION-400M/2B/5B) and CLIP (WiT).
• Specialized Architectures: CORnet (recurrent vision models), Segment Anything (SAM), and ALIGN.
• Pre-trained on Ecoset: Alexnet, VGG-16, Resnet50, and Inception_v3.
• Harmonization Models: Default is ViT_B16; also supports ResNet50, VGG16, EfficientNetB0, tiny_ConvNeXT, tiny_MaxViT, and LeViT_small.
• DreamSim Models: Default is open_clip_vitb32; also supports clip_vitb32, dino_vitb16, and an ensemble.

If you encounter GPU memory issues during feature extraction, you can use one of two strategies:

1. Reduce Batch Size: Decrease the batch_size parameter in your DataLoader instance.
2. Switch to CPU: Run extraction on the CPU instead of the GPU. This is slower but allows for larger batch sizes. Set the device parameter to 'cpu' in the get_extractor function.

Even if activations fit in RAM, you may encounter MemoryError when attempting to save them to disk. To resolve this, use the split_activations helper function to split the activation matrix into multiple files.

• Splitting: Call split_activations to divide the matrix into batches. The split parameter controls the number of files created (default is 10).
• Merging: To reconstruct the full activation matrix after loading, use the merge_activations function.

Example workflow:

# To save split files
split_activations(activations, PATH, split=10)

# To load and merge them back
activations = merge_activations(PATH)

Use CKA to compare image features from two different model architectures (at the same layer) or two different layers within the same architecture.

To use CKA, initialize a CKA object using get_cka and then call the .compare() method with your feature sets.

from thingsvision.core.cka import get_cka

backend = "torch" # 'torch' or 'numpy'
m = # number of images (e.g., features_i.shape[0])
kernel = 'linear' # 'linear' or 'rbf'
unbiased = True # whether to compute an unbiased version
device = "cuda" # only for 'torch' backend; NumPy runs on CPU
sigma = None # required if kernel='rbf'

cka = get_cka(backend=backend, m=m, kernel=kernel, unbiased=unbiased, device=device, sigma=sigma)
rho = cka.compare(X=features_i, Y=features_j)

By default, thingsvision stores all extracted features in RAM. For large datasets, this may cause your system to run out of memory. To prevent this, write features directly to disk as they are extracted:

• Set the output_dir parameter in the extract_features function to your desired directory.
• When output_dir is provided, extract_features returns None instead of the feature matrix.
• Use the step_size parameter to control how many batches are processed before writing to disk (the default is tuned to use approximately 8GB of RAM).

# get extractor and dataloader 
extractor = ...
batches = ...

output_dir = '/path/to/output/directory'
extractor.extract_features(
    batches=batches,
    module_name=...,
    flatten_acts=True,
    output_dir=output_dir
) # returns None if output_dir is set

To use DreamSim, you must first install the dreamsim package:

$ pip install dreamsim==0.1.2

Set model_name='DreamSim' and source='custom'. Use model_parameters={'variant': '...'} to select one of the four available models.

Available variants:

• clip_vitb32
• open_clip_vitb32
• dino_vitb16 (Note: requires module_name='model.mlp')
• ensemble (Note: requires module_name='model.mlp')

Caution: For dino_vitb16 and ensemble, features can only be extracted from the model.mlp module, not the model block.

import torch
from thingsvision import get_extractor

extractor = get_extractor(
  model_name='DreamSim',
  module_name='model.mlp',
  source='custom',
  device='cuda' if torch.cuda.is_available() else 'cpu',
  pretrained=True,
  model_parameters={'variant': 'open_clip_vitb32'}
)

To extract features from Harmonization models, install the required dependencies (works best on Ubuntu; macOS support is currently limited):

$ pip install git+https://github.com/serre-lab/Harmonization.git
$ pip install keras-cv-attention-models>=1.3.5

Set model_name='Harmonization' and source='custom'. Specify the architecture via model_parameters={'variant': '...'}.

Available variants:

• ViT_B16
• ResNet50
• VGG16
• EfficientNetB0
• tiny_ConvNeXT
• tiny_MaxViT
• LeViT_small

import torch
from thingsvision import get_extractor

extractor = get_extractor(
  model_name='Harmonization',
  source='custom',
  device='cuda' if torch.cuda.is_available() else 'cpu',
  pretrained=True,
  model_parameters={'variant': 'ViT_B16'}
)