AttentionVisualizer

AttentionVisualizer#

class stable_pretraining.callbacks.AttentionVisualizer(name: str, attention: str, image: str = 'image', num_prefix_tokens: int = 1, cls_index: int = 0, threshold: float | None = 0.6, head_reduction: str | int = 'mean', grid_size: int | Tuple[int, int] | None = None, 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, cmap: str = 'inferno', mask_color: Sequence[float] = (1.0, 0.15, 0.15), mask_alpha: float = 0.45, upsample_mode: str = 'bilinear', output_key: str | None = None, cell_size: float = 2.2, dpi: int = 130, verbose: bool = None)[source]#

Bases: Callback

Visualise and threshold ViT CLS self-attention, DINO-style.

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

  1. reads attention ((B, heads, T, T) full matrix or (B, heads, T) CLS row) and image ((B, C, H, W)) from the batch/outputs dict,

  2. extracts the CLS→patch attention, dropping num_prefix_tokens prefix columns, and reshapes it onto the (gh, gw) patch grid,

  3. optionally builds a binary object mask by keeping the most-attended patches whose cumulative mass reaches threshold (DINO recipe),

  4. writes the upsampled mean attention map to batch[output_key] and, if thresholding, the mask to batch[f"{output_key}_mask"], and

  5. renders a grid (input + per-head or mean heat maps + mask overlay) and logs it under {stage}/{name}.

Parameters:

  • name – Unique identifier — log tag and default dict-key prefix.

  • attention – Batch/outputs key holding attention weights.

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

  • num_prefix_tokens – Prefix tokens (CLS + registers) in the key/query dimension. The CLS row is read at cls_index; these columns are stripped before reshaping to the grid.

  • cls_index – Row index of the CLS token in a full attention matrix.

  • threshold – Cumulative-mass fraction in (0, 1) for the object mask; None disables the mask (heat maps only). DINO uses 0.6.

  • head_reduction – "mean" (average heads), "all" (one column per head), or an int head index.

  • grid_size – Explicit (gh, gw) patch grid; inferred when None.

  • 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 – Normalisation stats to invert for display.

  • image_std – Normalisation stats to invert for display.

  • cmap – Matplotlib colormap for the attention heat maps.

  • mask_color – RGB overlay colour for the thresholded object mask.

  • mask_alpha – Overlay opacity in [0, 1].

  • upsample_mode – Interpolation for heat maps (masks always use nearest).

  • output_key – Dict-key prefix. 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.AttentionVisualizer(
...     name="cls_attn",
...     attention="last_selfattention",  # (B, heads, T, T)
...     image="image",
...     num_prefix_tokens=1,
...     threshold=0.6,  # DINO object mask
...     head_reduction="all",  # show every head
... )
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 AttentionVisualizer:#

DINO feature PCA & CLS-attention visualisation

DINO feature PCA & CLS-attention visualisation