PCATokenVisualizer

PCATokenVisualizer#

class stable_pretraining.callbacks.PCATokenVisualizer(name: str, features: str, image: str = 'image', num_prefix_tokens: int = 1, n_components: int = 3, grid_size: int | Tuple[int, int] | None = None, normalize_features: bool = True, per_image: bool = False, foreground_threshold: float | None = None, background_color: Sequence[float] = (0.0, 0.0, 0.0), quantile: float = 0.02, max_images: int = 8, every_n_epochs: int = 1, log_on: Sequence[str] = ('val',), image_mean: Sequence[float] | None = None, image_std: Sequence[float] | None = None, upsample_mode: str = 'nearest', output_key: str | None = None, cell_size: float = 2.2, dpi: int = 130, verbose: bool = None)[source]#

Bases: Callback

Visualise ViT patch tokens as an RGB PCA image, DINO-style.

The callback fires at the start (batch_idx == 0) of each configured stage, once every every_n_epochs epochs. On each firing it:

  1. reads features ((B, N, D) tokens, or (B, D, H, W) maps) and image ((B, C, H, W)) from the batch/outputs dict,

  2. drops the num_prefix_tokens prefix tokens (CLS/registers),

  3. fits a joint PCA over every patch and maps the leading components to RGB (optionally using the DINOv2 foreground/background split),

  4. writes the upsampled (B, 3, H, W) RGB tensor to batch[output_key] (output_key defaults to name), and

  5. renders a [original | PCA] grid for up to max_images samples and logs it under {stage}/{name} via any media-capable logger.

Parameters:

  • name – Unique identifier — used as the log tag and the default dict key.

  • features – Batch/outputs key holding patch tokens (B, N, D) or a conv feature map (B, D, H, W).

  • image – Batch/outputs key holding the raw images (B, C, H, W).

  • num_prefix_tokens – Prefix tokens (CLS + registers) to strip from token inputs. Ignored for (B, D, H, W) inputs.

  • n_components – PCA components mapped to colour (3 → RGB).

  • grid_size – Explicit (gh, gw) patch grid. Inferred from the token count / image aspect ratio when None.

  • normalize_features – L2-normalise tokens before PCA (recommended).

  • per_image – Fit the PCA basis separately per image (True) or jointly over the whole batch (False, default). Joint keeps colours comparable across images (same material → same colour); per-image gives each sample its own basis and maximal colour contrast.

  • foreground_threshold – If set (0–1), reproduce the DINOv2 foreground trick — background patches (low first component) are painted with background_color and the object is coloured by a second PCA fit on foreground patches only.

  • background_color – RGB in [0, 1] for background patches.

  • quantile – Tail fraction clipped when scaling components to colour.

  • max_images – Max samples drawn in the figure.

  • every_n_epochs – Epoch interval between figures.

  • log_on – Stages to fire on — any of "train", "val", "test".

  • image_mean – Per-channel normalisation stats to invert for display. When None images are min-max normalised per sample.

  • image_std – Per-channel normalisation stats to invert for display. When None images are min-max normalised per sample.

  • upsample_mode – Interpolation used to bring the patch grid up to image resolution ("nearest" keeps crisp patch boundaries).

  • output_key – Dict key for the RGB tensor. Defaults to name.

  • cell_size – Figure geometry.

  • dpi – Figure geometry.

  • verbose – Verbosity; None derives from the global log level.

Example

>>> import stable_pretraining as spt
>>> viz = spt.callbacks.PCATokenVisualizer(
...     name="pca",
...     features="patch_tokens",  # (B, N, D) from a ViT
...     image="image",
...     num_prefix_tokens=1,  # strip CLS
...     foreground_threshold=0.6,  # DINOv2 look
... )
on_test_batch_end(trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=0)[source]#

Called when the test batch ends.

on_train_batch_end(trainer, pl_module, outputs, batch, batch_idx)[source]#

Called when the train batch ends.

Note

The value outputs["loss"] here will be the normalized value w.r.t accumulate_grad_batches of the loss returned from training_step.

on_validation_batch_end(trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=0)[source]#

Called when the validation batch ends.

Examples using PCATokenVisualizer:#

DINO feature PCA & CLS-attention visualisation

DINO feature PCA & CLS-attention visualisation