PaDISSolver.py

Source: LION/optimizers/PaDISSolver.py

LION-native solver for PaDIS patch-prior training.

class LION.optimizers.PaDISSolver.PaDISSolver(model, optimizer, loss_fn, geometry=None, verbose=True, device=None, solver_params=None, save_folder=None)

Bases: LIONsolver

Train a patch or whole-image PaDIS diffusion prior.

The solver adds patch scheduling, optional absolute-position channels, gradient-accumulation microbatches, patch-count budgets, exponential moving averages, validation-intensive selection, and resumable full-state checkpoints to LIONsolver.

Parameters:

• model (LIONmodel) – NCSN++ denoiser to optimise.

• optimizer (torch.optim.Optimizer) – Parameter optimiser owned by the solver.

• loss_fn (callable) – PaDIS-compatible denoising loss.

• geometry (Geometry, optional) – Image geometry used by the base solver.

• verbose (bool, optional) – Display progress information.

• device (torch.device, optional) – Training device.

• solver_params (SolverParams, optional) – Settings returned by default_parameters().

• save_folder (pathlib.Path, optional) – Root for checkpoints, sidecars, and final state.

static default_parameters(mode='padis-paper-ct-256')

Return training parameters for a PaDIS prior preset.

Parameters:

mode (str, optional) – Patch/whole-image preset. Patch ablations and -no-position variants are supported.

Returns:

Patch schedule, EDM bounds, EMA policy, and input convention.

Return type:

SolverParams

mini_batch_step(sino_batch, target_batch, patch_size=None)

Evaluate one PaDIS denoising mini-batch without stepping the optimiser.

Parameters:

patch_size (int | None)

Return type:

Tensor

train_step()

Train for one complete pass through the configured loader.

train_for_patches(target_patches, *, validation_interval_patches=None, validation_max_patches=None, validation_repeat_until_max_patches=False, checkpoint_interval_patches=None, checkpoint_interval_seconds=None, log_interval_patches=None, max_train_seconds=None, log_fn=None)

Train until seen_patches reaches target_patches.

This is the PaDIS-native budget: the underlying PaDIS repository uses an image/patch counter rather than epochs to drive LR ramp-up, EMA, and run duration.

Parameters:

• target_patches (int)

• validation_interval_patches (int | None)

• validation_max_patches (int | None)

• validation_repeat_until_max_patches (bool)

• checkpoint_interval_patches (int | None)

• checkpoint_interval_seconds (float | None)

• log_interval_patches (int | None)

• max_train_seconds (float | None)

• log_fn (Callable[[dict[str, object], int], None] | None)

Return type:

None

validate(max_patches=None, *, repeat_until_max_patches=False)

Evaluate EMA validation loss over a bounded patch sample.

Parameters:

• max_patches (int, optional) – Maximum number of image patches evaluated.

• repeat_until_max_patches (bool, optional) – Restart a finite loader until max_patches is reached.

Returns:

Patch-weighted mean validation loss.

Return type:

float

set_checkpoint_retention(max_periodic_checkpoints)

Configure how many ordinary periodic checkpoints are retained.

Parameters:

max_periodic_checkpoints (int | None)

Return type:

None

prune_periodic_checkpoints(max_periodic_checkpoints=None)

Remove the oldest periodic checkpoints and their sidecars.

Parameters:

max_periodic_checkpoints (int | None)

Return type:

None

save_checkpoint(epoch)

Save model, optimiser, EMA, and patch-counter checkpoint state.

save_validation(epoch)

Save validation-selected EMA weights and resumable state.

save_final_results(final_result_fname=None, save_folder=None, epoch=None)

Save final EMA model weights and a full training-state sidecar.

load_checkpoint()

Restore the newest compatible periodic or full-state checkpoint.