peal.architectures.predictors.TorchvisionModel

class peal.architectures.predictors.TorchvisionModel

Bases: Module

Pretrained backbone plus a fresh linear head, behind one interface.

The model string selects the backbone and decides where the classification head lives: for the torchvision ResNets and ViTs the head inside the backbone is replaced, for the feature-extractor backbones (DINOv2/DINOv3, OpenCLIP, UNI) a separate self.fc is added on top of the frozen-format embedding. All variants are reachable through feature_extractor(), get_last_layer() and forward(), which is what the PEAL adaptors rely on.

Recognised model values

resnet18, resnet50

Pretrained torchvision ResNet with a bias-free fc.

resnet_loaded

A pickled model given by config.base_model, whose inner fc is replaced.

dino_v2/dino_v2_small/dino_v2_base

Hugging Face DINOv2 large/small/base plus its image processor.

dino_v3/dino_v3_small/dino_v3_base/dino_v3:<repo>

DINOv3 from Hugging Face, falling back to the public timm weights when the repo is gated (self.is_timm records which path was taken).

open_clip:<model>:<pretrained>

OpenCLIP image tower; the normalisation stats and input size are read back out of the generated preprocessing transform.

PLIP, UNI

Pathology foundation models.

vit_b_16

torchvision ViT; the positional embedding is truncated or re-inited when input_size differs from 224.

param model:

Backbone selector, see above.

type model:

str

param num_classes:

Number of output logits of the classification head.

type num_classes:

int

param input_size:

Square input resolution; only used to adapt the ViT positional embedding.

type input_size:

int, optional

param config:

Predictor config; only config.base_model is read, for resnet_loaded.

type config:

object, optional

raises ValueError:

If model is not one of the supported names.

Methods

feature_extractor	Embed a batch of images with the backbone, without the head.
forward	Classify a batch, optionally also returning the latent code.
get_last_layer	Return the final linear layer, i.e. the classification head.

__init__(model, num_classes, input_size=None, config=None)

Build the selected backbone and attach the classification head.

feature_extractor(x)

Embed a batch of images with the backbone, without the head.

The path taken depends on self.model_type:

• ResNets (and models predating the model_type attribute): every child module but the last is re-assembled into a Sequential, so the result still carries the trailing spatial dimensions (B, C, 1, 1).

• DINOv2/DINOv3: the batch is resized to the processor’s crop size (224 for the timm fallback), normalised with the processor’s image_mean/image_std and reduced to the CLS token.

• OpenCLIP and PLIP: resized to self.clip_size, normalised with the CLIP stats and passed through the image tower.

• UNI: the timm transform is applied and the model called directly.

• otherwise (torchvision ViT): patch embedding, CLS token, encoder, CLS output.

Parameters:

x (torch.Tensor) – Images of shape (B, C, H, W) in [0, 1].

Returns:

Latent codes, (B, D) for the transformer backbones and (B, D, 1, 1) for the ResNets.

Return type:

torch.Tensor

get_last_layer()

Return the final linear layer, i.e. the classification head.

This is self.model.fc for the ResNets and the separate self.fc for the DINO, OpenCLIP and UNI backbones. Adaptors such as CLARC and the projection adaptor edit the weights of this layer.

Returns:

None for backbones whose head is not exposed this way (for example the torchvision ViT).

Return type:

torch.nn.Linear or None

forward(x, return_latents=False)

Classify a batch, optionally also returning the latent code.

Parameters:

• x (torch.Tensor) – Images of shape (B, C, H, W).

• return_latents (bool, optional) – Return (latent_code, logits) instead of just the logits. For the ResNets the latent code is squeezed to (B, D); for the torchvision ViT branch the logits are returned twice.

Returns:

Logits of shape (B, num_classes), or the pair described above.

Return type:

torch.Tensor or tuple of torch.Tensor