LION.optimizers.GaussianDenoiserSolver

Classes

GaussianDenoiserSolver(model, optimizer, loss_fn)

Trains a Gaussian denoiser (e.g. for Plug-and-Play): clean images (the targets), with Gaussian noise added, are its inputs.
class LION.optimizers.GaussianDenoiserSolver.GaussianDenoiserSolver(model, optimizer, loss_fn, operator_params=None, solver_params=None, model_regularization=None, scheduler=None, scheduler_step='epoch', verbose=1, device=None, seed=0, deterministic=False)[source]

Bases: LIONsolver

Trains a Gaussian denoiser (e.g. for Plug-and-Play): clean images (the targets), with Gaussian noise added, are its inputs. As in the literature (DnCNN, FFDNet and DRUNet in KAIR, Zhang et al.; the gradient-step denoiser, Hurault et al.), each image gets its own noise level, drawn from a range, and is randomly flipped and rotated.

Noise levels are in the units of the data the model sees: the noise is added after the model’s normalisation (model.set_normalisation()), if any. The literature, and LION, assume images in [0, 1] (e.g. a “dataset” or “custom” normalisation to [0, 1]), where its noise levels are given as k/255: 25/255 = 0.098, 50/255 = 0.196. With another range (e.g. no normalisation), scale the noise levels with it. The same holds for the noise level a PnP reconstruction gives the denoiser.

solver_params (default_parameters()):

  • noise_level: the noise’s standard deviation: one value, or a range [low, high], drawn uniformly for each image (or patch). A model that takes the noise level (model.takes_noise_level()) is given one per image.

  • patches: None (whole images), or n_patches random crops of patch_size x patch_size of each image (GaussianDenoiserSolver.patches(n_patches, patch_size)).

  • augment: each training image (or patch) flipped and rotated by a random one of the 8 symmetries of the square (of the rectangle’s 4, for images that aren’t square).

  • random_erasing: each training image (or patch), half of the time, with a random rectangle set to 0 (image_transforms.random_erasing).

  • validation_noise_levels: None, validation draws the noise levels as training does (the same ones every time); or a list of them (e.g. [15 / 255, 25 / 255, 50 / 255]), every validation image at each, the validation value being their mean.

For PnP methods that converge, a penalty on the denoiser’s Jacobian (LION.losses.jacobian_penalty) is the solver’s model_regularization: FirmlyNonexpansivePenalty() (Pesquet et al. 2021) or ProximalDenoiserPenalty() (Hurault et al. 2022, for a gradient-step denoiser). Both papers add it after training without it (e.g. fine-tuning with LIONmodel.load and another solver). Cite the one used (its cite()).

Every solver takes (model, optimizer, loss_fn, operator_params, solver_params, model_regularization, scheduler, scheduler_step, verbose, device, seed, deterministic), in that order (ARSolver has no loss_fn).

operator_params: the parameters of the operator (for CT, the geometry); the model’s, if it has them.

solver_params: the solver’s settings, from default_parameters().

model_regularization: None, or a term added to the training loss of every mini-batch, model_regularization(model, step): a penalty on the model’s weights (step unused) or on what it does with the mini-batch (e.g. LION.losses.jacobian_penalty). step has input (what the model was given), output, target (what the output is compared with; None for solvers without one) and forward (the model as the solver calls it, e.g. with the noise level). Recorded as loss_fn is.

scheduler: a torch learning-rate scheduler made on optimizer (e.g. torch.optim.lr_scheduler.StepLR(optimizer, 30)), or None. It steps after every epoch (scheduler_step="epoch", torch’s convention) or every mini-batch ("batch", e.g. for OneCycleLR). ReduceLROnPlateau steps with the validation value of the epochs that validate, or with the training loss when there is no validation. Its state is saved in checkpoints, and the training record has it and the learning rate of every epoch.

seed: the training seed. When training starts, torch (CPU and CUDA), numpy and Python’s random are seeded with it, so the same solver, model and data give the same training, and checkpoints keep the random state, so a resumed run continues the same way. The randomness within a run (shuffling, noise a solver adds) still changes every batch and epoch. Two runs with the same seed give the same model: use different seeds for an ensemble. None: unseeded.

deterministic: use only deterministic algorithms (torch’s use_deterministic_algorithms, deterministic cuDNN), so that on a GPU two runs with the same seed give exactly the same model, and a resumed run the uninterrupted one. Slower; a few operations then raise (torch names them). It applies to the whole program from here on. Off, GPU runs with the same seed differ slightly (e.g. ~1e-3 in the weights after a few epochs). On a GPU it needs CUBLAS_WORKSPACE_CONFIG=:4096:8 before cuBLAS starts: it is set here if it isn’t (exporting it before the script is surer).

verbose: how much progress is shown. 0: nothing; 1 (default): one progress bar for training (every mini-batch of every epoch, with the epoch, the loss, the last validation value and the time left) and one for testing; 2: also the validation bars and messages (each epoch’s losses, checkpoints, the seed). True is 1, False 0.

mini_batch_step(y, target)[source]

The loss of the model on target (clean images) with noise added (y isn’t used: the inputs are made from the targets).

static default_parameters()[source]

The settings in the class docstring, with their defaults.

static patches(n_patches, patch_size)[source]

solver_params.patches: n_patches random crops of patch_size x patch_size of each image.

validate()[source]

The validation value: validation_fn of the model’s outputs on the validation images with noise added, against the images (on patches, if training is; not flipped, rotated or erased). The noise is the same in every validation.