prior_target_standardization

prior_target_standardization

Callback for fitting learned-prior latent target standardization.

PriorTargetStandardizationCallback

PriorTargetStandardizationCallback(max_batches=None, require_posterior_frozen=True)

Bases: Callback

Fit latent target statistics for learned denoising priors.

Initialize the callback.

Parameters:

Name	Type	Description	Default
max_batches	int | None	Maximum train batches to scan when fitting stats. None scans the full train dataloader.	None
require_posterior_frozen	bool	If true, fit only after the posterior encoder has no trainable parameters.	True

Raises:

Type	Description
ValueError	If max_batches is not positive when provided.

Source code in src/versatil/training/callbacks/prior_target_standardization.py

def __init__(
    self,
    max_batches: int | None = None,
    require_posterior_frozen: bool = True,
) -> None:
    """Initialize the callback.

    Args:
        max_batches: Maximum train batches to scan when fitting stats. ``None``
            scans the full train dataloader.
        require_posterior_frozen: If true, fit only after the posterior encoder
            has no trainable parameters.

    Raises:
        ValueError: If ``max_batches`` is not positive when provided.
    """
    super().__init__()
    if max_batches is not None and max_batches <= 0:
        raise ValueError(
            f"PriorTargetStandardizationCallback max_batches must be positive, "
            f"got {max_batches}."
        )
    self.max_batches = max_batches
    self.require_posterior_frozen = require_posterior_frozen

on_train_start

on_train_start(trainer, pl_module)

Fit stats if the initial training regime is already prior-only.

Parameters:

Name	Type	Description	Default
trainer	Trainer	Active Lightning trainer.	required
pl_module	LightningModule	Lightning module wrapping the policy.	required

Source code in src/versatil/training/callbacks/prior_target_standardization.py

def on_train_start(
    self, trainer: pl.Trainer, pl_module: pl.LightningModule
) -> None:
    """Fit stats if the initial training regime is already prior-only.

    Args:
        trainer: Active Lightning trainer.
        pl_module: Lightning module wrapping the policy.
    """
    self._fit_if_ready(trainer=trainer, pl_module=pl_module)

on_train_epoch_start

on_train_epoch_start(trainer, pl_module)

Fit stats when staged training first makes the prior trainable.

Parameters:

Name	Type	Description	Default
trainer	Trainer	Active Lightning trainer.	required
pl_module	LightningModule	Lightning module wrapping the policy.	required

Source code in src/versatil/training/callbacks/prior_target_standardization.py

def on_train_epoch_start(
    self, trainer: pl.Trainer, pl_module: pl.LightningModule
) -> None:
    """Fit stats when staged training first makes the prior trainable.

    Args:
        trainer: Active Lightning trainer.
        pl_module: Lightning module wrapping the policy.
    """
    self._fit_if_ready(trainer=trainer, pl_module=pl_module)