OnlineImageDecoder

class stable_pretraining.callbacks.OnlineImageDecoder(module: LightningModule, name: str, input: str, target: str, image_shape: Tuple[int, int, int], embed_dim: int, kind: str = 'auto', patch_size: int | None = None, decoder_kwargs: dict | None = None, loss: callable | None = None, metrics: dict | tuple | list | Metric | None = None, optimizer: str | dict | partial | Optimizer | None = None, scheduler: str | dict | partial | LRScheduler | None = None, accumulate_grad_batches: int = 1, gradient_clip_val: float = None, gradient_clip_algorithm: str = 'norm', verbose: bool = None)

Bases: OnlineProbe

Online probe that decodes embeddings back to images.

The decoder is constructed once at init time from image_shape, embed_dim, and (for ViT inputs) patch_size. From there the callback behaves exactly like OnlineProbe: features under input are detached and passed through the decoder, MSE is computed against target, and the optimizer/scheduler are managed independently of the main model.

The reconstruction is written to outputs[f"{name}_preds"] on every forward pass, so other callbacks (visualisers, video writers, custom metrics) can read it directly via get_data_from_batch_or_outputs(). For example, OnlineImageDecoder(name="recon", ...) exposes the reconstructed image at the key "recon_preds".

Parameters:

• module – The LightningModule being trained.

• name – Unique identifier (used for logging, prediction key, metric namespace, and stored under pl_module.callbacks_modules[name]).

• input – Key in batch or outputs holding the embedding. (B, D) for CNN decoding, (B, P, D) for ViT decoding.

• target – Key in batch or outputs holding the target image, shape (B, C, H, W) matching image_shape.

• image_shape – (channels, height, width) of the reconstruction target. Currently height == width (square images only).

• embed_dim – Feature dimension D of the input embedding.

• kind – "auto" (default), "cnn", or "vit". "auto" resolves to "vit" if patch_size is set, otherwise "cnn".

• patch_size – Patch side length for the ViT decoder. Must divide image_shape[1].

• decoder_kwargs – Extra kwargs forwarded to the underlying decoder (e.g. base_channels, decoder_dim, depth).

• loss – Reconstruction loss. Defaults to MSE with a shape-mismatch guard.

• metrics – Defaults to torchmetrics.MeanSquaredError.

• optimizer

• scheduler

• accumulate_grad_batches

• gradient_clip_val

:param : :param gradient_clip_algorithm: Forwarded to OnlineProbe. Optimizer defaults to

Adam(lr=1e-3, weight_decay=0) (a more sensible choice for a conv/transformer decoder than the probe’s default LARS).

Parameters:

• verbose – Forwarded to OnlineProbe. Optimizer defaults to Adam(lr=1e-3, weight_decay=0) (a more sensible choice for a conv/transformer decoder than the probe’s default LARS).

• Example

• -------

• image:: (Decode a pooled vector to a 64x64 RGB) –

  spt.callbacks.OnlineImageDecoder(

  module=module, name=”recon_rgb”, input=”embedding”, target=”image”, image_shape=(3, 64, 64), embed_dim=768,

  )

• map:: (Decode a token grid to a 96x96 depth) –

  spt.callbacks.OnlineImageDecoder(

  module=module, name=”recon_depth”, input=”patch_tokens”, target=”depth”, image_shape=(1, 96, 96), embed_dim=384, patch_size=8, # triggers the ViT decoder

  )