LION.optimizers.Equivariance2InverseSolver

Classes

Equivariance2InverseSolver(model, optimizer, ...)

Equivariance2Inverse (E2I).

Functions

estimate_noise_std(sinos[, noise_kernel])

Estimate the standard deviation of additive Gaussian noise in sinograms.

---

class LION.optimizers.Equivariance2InverseSolver.Equivariance2InverseSolver(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)

Bases: LIONsolver

Equivariance2Inverse (E2I).

The training loss is the mean squared error, as in the method’s formulation: a projection-wise cross-validation (data consistency) term plus equivariance_strength (lambda in the paper) times the equivariance term. loss_fn is only used as the default validation and testing function.

The equivariance term rotates the reconstruction by a random angle, projects it and adds new Gaussian noise to the (post-log) sinogram, as in the paper and the authors’ code (github.com/D1rk123/equivariance2inverse): iid with standard deviation noise_std or, with noise_kernel, blurred along the detector with that kernel (the detector’s point response). noise_std=None estimates it from the first training batch (estimate_noise_std) and stores it in the solver parameters.

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.

static default_parameters()

equivariance_strength (lambda), noise_std, noise_kernel and recon_fn, as below.

mini_batch_step(y, targets)

The loss: projection-wise cross-validation (one angle left out, the model’s output projected onto it) plus equivariance_strength times the equivariance term (targets aren’t used: self-supervised).

reconstruct(y)

The model’s reconstructions of the measurements y: of the reconstructions recon_fn makes of them, as in training.

static cite(cite_format='bib')

Where Equivariance2Inverse comes from: its citation, BibTeX (“bib”) or “MLA”.

LION.optimizers.Equivariance2InverseSolver.estimate_noise_std(sinos, noise_kernel=None)

Estimate the standard deviation of additive Gaussian noise in sinograms.

The noise is noise_kernel * w along the detector, with w iid Gaussian of the returned standard deviation (noise_kernel=None: iid noise). The estimate uses second differences along the detector, y[j-1] - 2 y[j] + y[j+1], where a smooth signal cancels: for that noise their standard deviation is std * ||second difference of noise_kernel||. A robust standard deviation (from the median absolute deviation) ignores the large differences at edges of the object. Triplets with a 0, e.g. post-log values clipped at 0 in air, are left out.

The noise of CT data depends on the signal (it is larger behind more attenuation), so this is a typical level over the sinograms. The E2I paper measures it in a background region instead; give the solver’s noise_std to use such a value.

sinos: sinograms (…, angles, detectors); noise_kernel: the 1D blur of the noise along the detector, or None. Returns the standard deviation before the blur.