LION.optimizers.LIONsolver

Classes

LIONsolver(model, optimizer, loss_fn[, ...])

Base class for the solvers, which train a model.

Functions

model_forward(model, data, op)	Run model on a batch of measurements, as the solvers do (and LION.reconstructors.ForwardReconstructor).
warn_if_not_mse(loss_fn, solver_name)	Warn that loss_fn is not used by a solver whose loss must be the MSE.

---

class LION.optimizers.LIONsolver.LIONsolver(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: ABC

Base class for the solvers, which train a model.

A solver is made with the model, the optimizer, the loss and its settings (__init__), then told about the data and what to do around training, then trained:

• set_training(loader): the training data (required);

• set_validation(loader, freq, fn): validate every freq epochs, keeping the best model;

• set_saving(folder, name): where the final and best models go;

• set_checkpointing(name, freq): checkpoints, to resume from;

• set_testing(loader, fn): data for test();

• train(n_epochs), then save_final_results().

Every saved file has the training record (training_record): enough to train the same model again (LIONsolver.instantiate, LION.package).

A new solver derives from this class and defines mini_batch_step(y, target) (the loss of one mini-batch) and default_parameters() (its solver_params); see SupervisedSolver for the simplest.

operator_params: the parameters of the operator the measurements come from (an OperatorParameters; for CT, the geometry). Not needed when the model has them: the solver takes the model’s.

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.

property do_normalise

True if the model has a normalisation set (model.set_normalisation()).

reconstruct_measurements(y, op)

Reconstruct the measurements y of the operator op with the solver parameters’ recon_fn(y, op), for solvers whose method starts from reconstructions (Noise2Inverse, Proj2Proj, …). Its default, "default", is the operator’s default reconstruction (FDK for CT), replaced by that function when the solver is made.

spec_kwargs()

The solver’s own constructor arguments: not the model (saved as the model) nor the optimizer (built on the model’s parameters).

training_record(kind, epoch, value=None)

The record of this training, for a file holding the kind (“final”, “best_validation”, “checkpoint”) weights after epoch epochs (see LION.optimizers.training_record).

classmethod instantiate(source, model, *, allow=(), **arguments)

The solver a saved file describes, training model.

source: a model file a solver saved (or the dict read from it), or its training record. Made from the solver’s own record only: its class, solver_params, loss and seed, the optimizer (class, hyperparameters and parameter groups) on model’s parameters, and the learning-rate scheduler on that optimizer. The data, validation, testing and saving are set as usual (set_training, …). arguments replace recorded ones or give those that couldn’t be recorded (e.g. loss_fn for a lambda loss) and where it runs (device, verbose). Called on LIONsolver it makes the class the record names; on a solver class, that class.

classmethod parameters_from_file(source, *, allow=())

What default_parameters() gives, as source (a model file a solver saved, or its training record) recorded it: the solver’s solver_params, to change and make the solver with (instantiate(source, model, solver_params=parameters)). The defaults for a solver that takes none (e.g. SupervisedSolver).

abstract static default_parameters()

The solver’s settings (solver_params), each with its default: every setting the solver has, so this is also their list.

set_training(train_loader, loss_fn=None)

The training data: a DataLoader of (measurements, target) batches. loss_fn: replaces the solver’s loss, if given.

set_validation(validation_loader, validation_freq, validation_fn=None, validation_fname=None, reconstructor=None)

Validate every validation_freq epochs on validation_loader, with validation_fn(output, target) (the training loss if None).

The model with the best validation value is saved as validation_fname in the solver’s save folder (see set_saving()). Best is the lowest value, unless validation_fn has higher_is_better = True (as LION’s PSNR and SSIM metrics do; any function can set it, e.g. my_metric.higher_is_better = True).

reconstructor: a reconstructor (LION.reconstructors) to reconstruct the validation data with, using the model being trained (set as its model); otherwise the solver’s own _forward. It is made outside the solver, with its own parameters, and recorded.

set_testing(test_loader, testing_fn=None, reconstructor=None)

The data test() uses: test_loader, scored with testing_fn(output, target) (the training loss if None). reconstructor: as in set_validation.

set_saving(save_folder, final_result_fname)

Where results are saved: the final model as save_folder/final_result_fname (save_final_results()), and the best validation model next to it. save_folder must be a folder that exists (ValueError otherwise).

set_checkpointing(checkpoint_fname, checkpoint_freq=10, load_checkpoint_if_exists=True, save_folder=None)

Save a checkpoint every checkpoint_freq epochs, as checkpoint_fname (its * replaced by the epoch) in save_folder (the solver’s, from set_saving(), if None). A checkpoint has the model, the optimizer’s and scheduler’s state, the losses and the random state: with load_checkpoint_if_exists, train() resumes from the latest one, as if never stopped.

check_training_ready(error=True, autofill=True)

Checks what training needs (all set when the solver is made).

error: raise when something required is missing (otherwise warn); autofill: fill in what has a default (the device). Returns 0 when ready, 1 when something is missing, 2 when something has the wrong type.

check_validation_ready(autofill=True)

Checks what validation needs (set_validation()), filling in the validation function (the loss) and frequency (10) with autofill, and the best model’s file name. Returns 0 when ready, 1 when something is missing, 2 when something has the wrong type.

check_testing_ready(error=True)

Checks what test() needs (set_testing()); error: raise when something is missing (otherwise warn). Returns as check_validation_ready.

check_checkpointing_ready(autofill=True)

Fills in what checkpointing needs; checkpointing is optional, so nothing missing here is warned about.

check_saving_ready()

Checks that there is a folder to save in (set_saving()): 0 if so, 1 if not set, 2 if not a folder.

check_complete(error=True, autofill=True)

Every check (training, testing, saving, validation, checkpointing): the worst of their return codes.

save_checkpoint(epoch)

Saves a checkpoint after epoch (0-based), as set_checkpointing() says: the model, the optimizer’s and scheduler’s state, the losses, the random state and the training record.

save_validation(epoch)

Saves the model as it is after epoch, its best validation so far, as validation_fname in the save folder, with the training record and the validation value.

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

Saves the trained model as final_result_fname in save_folder (those of set_saving() if None), with the training record and the losses. epoch: the epochs trained (the current one if None).

clean_checkpoints()

Deletes the checkpoint files (set_checkpointing()), e.g. after a run that finished.

test()

Scores the model on the test data (set_testing()): an array of testing_fn values, one per batch (or per sample, if testing_fn gives one per sample). Empty, with a warning, without test data.

load_checkpoint()

Loads the latest checkpoint, if there is one, and resumes from it: the model, the optimizer’s and scheduler’s state, the epoch, the losses, the random state. train() calls it.

train_step()

One pass over the training data: for each mini-batch, the loss (mini_batch_step), a gradient step (and a scheduler step, with scheduler_step="batch"). Returns the mean loss of the epoch.

epoch_step(epoch)

Epoch epoch (0-based): training (train_step), validation when due (saving the best model), and a scheduler step.

train(n_epochs)

Trains until n_epochs epochs are done in all: resuming from a checkpoint (set_checkpointing()) or an earlier train(), with one progress bar (verbose). The model is left in training mode; save_final_results() saves it.

validate()

The validation value of the model now: the mean of validation_fn over the validation data, with the model in evaluation mode, the random state left as it was.

get_model()

The model being trained.

abstract mini_batch_step(y, target_batch)

The loss of one mini-batch: measurements y and their targets (the solver adds model_regularization with _regularized). What a new solver defines.

static cite(cite_format='bib')

Where the solver comes from: its paper’s citation, BibTeX (“bib”) or “MLA” (LION.utils.citation); None if it has none. Solvers from a paper define it.

LION.optimizers.LIONsolver.model_forward(model, data, op)

Run model on a batch of measurements, as the solvers do (and LION.reconstructors.ForwardReconstructor).

Image-domain models get a reconstruction (with operator op) as input: the operator’s default (FDK for CT), or the model’s own (model.initial_reconstruction). Sinogram-domain models get the measurements. With a normalisation set (model.set_normalisation()), the input is normalised (for measurements: as the images are, by linearity) and the output mapped back, so the output is in the units of the target.

LION.optimizers.LIONsolver.warn_if_not_mse(loss_fn, solver_name)

Warn that loss_fn is not used by a solver whose loss must be the MSE.

Some self-supervised losses are only valid for the mean squared error (e.g. their target is a conditional mean), so these solvers use the MSE whatever loss is passed; loss_fn is still the default validation and testing function.