ReconViz

class stable_pretraining.callbacks.ReconViz(specs: Sequence[Tuple[str, str, str | None]], mode: str = 'grid', pixel_mean: Sequence[float] | Tensor | None = None, pixel_std: Sequence[float] | Tensor | None = None, n_items: int = 8, seq_len: int | None = None, fps: int = 2, verbose: bool = None)

Bases: Callback

Render decoder reconstructions as target | recon media each val epoch.

Pairs one-to-one with OnlineImageDecoder: point a spec at the decoder’s name and ReconViz reads its f"{name}_preds" output, stitches each reconstruction next to its target, and logs the result.

On on_validation_batch_end (only batch_idx == 0) the first N rows of predictions and their matching targets are cached (detached, float, CPU). On on_validation_epoch_end the cache is denormalised, clamped to [0, 1], stitched (target | separator | recon) and emitted, then cleared.

Parameters:

• specs – List of (preds_key_or_name, target_key, caption) tuples. The first element is resolved against outputs/batch by trying f"{key}_preds" first (the OnlineImageDecoder convention) then key verbatim, so both "recon" and "recon_preds" work. caption may be None.

• mode – "grid" (default) treats the cached rows as a plain (B, C, H, W) batch and logs a grid image via log_image. "video" treats them as a trajectory-contiguous flat layout (n_items trajectories × seq_len transitions) and logs an MP4 via log_video.

• pixel_mean – Per-channel mean used at normalisation time, broadcastable to (C, 1, 1) (a length-C sequence is accepted and reshaped). Reconstruction is x * pixel_std + pixel_mean. Must be given together with pixel_std; when both are None denorm is a no-op.

• pixel_std – Per-channel std used at normalisation time (see pixel_mean).

• n_items – "grid" — number of images shown. "video" — number of trajectories (n_traj) shown side by side per frame.

• seq_len – "video" only — transitions per trajectory. Required when mode="video"; the flat layout is assumed to be n_items * seq_len contiguous rows.

• fps – "video" only — frames per second of the emitted MP4.

• verbose – Enable per-emit info logging. None defers to the global verbosity setting.

Example

Pair with OnlineImageDecoder to log a target | recon grid:

dec = spt.callbacks.OnlineImageDecoder(
    module=module,
    name="recon",
    input="embedding",
    target="image",
    image_shape=(3, 64, 64),
    embed_dim=768,
)
viz = spt.callbacks.ReconViz(
    [("recon", "image", "recon")],
    mode="grid",
)

on_validation_batch_end(trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=0) → None

Called when the validation batch ends.

on_validation_epoch_end(trainer: Trainer, pl_module: LightningModule) → None

Called when the val epoch ends.