Training Models and Selecting Checkpoints

The training pipeline covers patch and whole-image diffusion priors, native 512-resolution training, fixed-noise and noise-conditioned PnP denoisers, validation-intensive final checkpoints, resumption, and W&B provenance.

Training models and checkpoints

The GCP runner’s default task order is:

Task	Model/data change
whole_lidc_full	Whole-image prior, all processed training slices.
whole_lidc_default	Whole-image prior, four-slice regime.
pnp_lidc_drunet	Unconditioned DRUNet.
pnp_lidc_drunet_noise_cond	DRUNet with a noise-level input channel.
patch_lidc_p96_default	Patch prior with maximum patch size 96.
patch_lidc_full	Default patch prior, all processed training slices.
patch_lidc_512	Native-resolution patch prior.
patch_lidc_default	Main 256-resolution patch prior.
patch_lidc_p32_default	Maximum patch size 32.
patch_lidc_p16_default	Maximum patch size 16.
patch_lidc_p8_default	Fixed patch size 8.
patch_lidc_no_pos_default	Main patch prior without x/y coordinate channels.

Diffusion training defaults

Patch and whole-image priors use an EDM-preconditioned NCSN++ denoiser. Both have four resolution levels, four residual blocks per level, 128 base channels, dropout 0.05, and attention at resolution 16. Patch channel multipliers are [1, 2, 2, 2]; whole-image multipliers are [1, 2, 2, 4]. Default models receive the image plus normalised x/y coordinates. The no-position model receives only the image channel.

Adam uses learning rate 2e-4. Training noise is log-normal, log(sigma) ~ N(-1.2, 1.2^2), truncated to [0.002, 40], with sigma_data=0.5; the learning rate is ramped over the first ten million training images or patches. Whole-image batches contain 8 slices. Patch effective batch sizes are 128, 256, and 512 from largest to smallest patch. An EMA with a 500,000-example half-life supplies validation and inference weights.

Training preset	Patch sizes	Probabilities	Padding
P=8	8	1.0	8
P=16	8, 16	0.3, 0.7	16
P=32	8, 16, 32	0.2, 0.3, 0.5	32
Default P=56	16, 32, 56	0.2, 0.3, 0.5	24
P=96	32, 64, 96	0.2, 0.3, 0.5	32
Native 512	16, 32, 64	0.2, 0.3, 0.5	64

The GCP base phase budgets six hours for each patch task and eighteen hours for each whole-image task. The following six-hour validation-intensive phase uses 4,000 validation patches every 20,000 seen patches, or all 328 validation images every 2,500 seen images. The best EMA checkpoint in that phase is named <prefix>_min_intense_val.pt and is the diffusion checkpoint used by the reconstruction matrix.

Slurm follows the same two-phase schedule within each job. Patch-prior tasks train normally for six hours and intensively validate for six hours; their Slurm allocation is 12 hours plus a 30-minute setup/final-save buffer. Whole-image tasks train normally for eighteen hours and intensively validate for the final six hours; their allocation is 24 hours plus the same buffer. The two task families are submitted as separate arrays so their scheduler time limits are accurate. P=96 remains capped at batch size 96 on A100 by default, even if the general patch batch size is raised; lower it with PADIS_P96_BATCH_SIZE or change the explicit safety ceiling with PADIS_P96_A100_BATCH_LIMIT only after a memory check.

PnP defaults

Both DRUNets use 64 internal channels, four blocks per level, batch size 8, Adam at 1e-4, and 100 epochs. Gaussian training noise is sampled uniformly from [0, 0.05]. The best validation files are:

pnp_lidc_drunet/pnp_lidc_drunet_min_val.pt
pnp_lidc_drunet_noise_cond/pnp_lidc_drunet_noise_cond_min_val.pt

The conditioned reconstruction row supplies noise level 0.03.

Resuming and W&B

Periodic full-state files contain optimiser, counters, and EMA state. Training scripts discover compatible periodic checkpoints automatically. The GCP runner keeps per-task phase markers under:

<training-root>/.gcp_spot/{done,running,failed,logs,runtime}/

Rerunning the same command skips valid completed phases and resumes incomplete ones. W&B logging defaults to online mode and final artefacts are uploaded unless PADIS_NO_WANDB_ARTIFACT=1; persistent checkpoints do not depend on W&B.

Notes and limitations

The cp_tv implementation executes LION’s Chambolle-Pock TV solver and is not the exact ADMM-TV algorithm described by Hu et al. The baseline uses fan-beam FDK, not parallel-beam FBP. pnp_admm uses a LION-native DRUNet surrogate as its DRUNet denoiser; the source description does not give enough optimizer, architecture, or stopping-rule detail for exact identity. A no-PaDIS-prior or empty checkpoint is never treated as checkpoint identity. The completed final training and reconstruction used a Google-hosted NVIDIA RTX PRO 6000; fixed-validation tuning also used a GTX 1070. The Slurm A100 path is an equivalent reproduction backend, not the hardware on which the reported final results were obtained.

The CP baseline retains the isotropic TV objective, uses lambda=0.001 for 1,000 iterations, and clips its output to [0,1] without imposing non-negativity inside the solver. It is therefore reported as CP, not as an exact reproduction of the original work’s ADMM-TV baseline.