LIONsolver.py

Source: LION/optimizers/LIONsolver.py

class LION.optimizers.LIONsolver.SolverState(value, names=_not_given, *values, module=None, qualname=None, type=None, start=1, boundary=None)

Bases: Enum

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

Bases: ABC

Parameters:

• model (LIONmodel)

• optimizer (Optimizer)

• loss_fn (Callable)

• geometry (Geometry)

• verbose (bool)

• device (torch.device)

• solver_params (Optional[SolverParams])

• save_folder (Optional[pathlib.Path])

set_training(train_loader, loss_fn=None)

This function sets the training data

Parameters:

• train_loader (DataLoader)

• loss_fn (Callable | None)

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

This function sets the validation data

Parameters:

• validation_loader (DataLoader)

• validation_freq (int)

• validation_fn (Callable | None)

• validation_fname (str | None)

• save_folder (Path | None)

set_testing(test_loader, testing_fn=None)

This function sets the testing data

Parameters:

• test_loader (DataLoader)

• testing_fn (Callable | None)

set_saving(save_folder, final_result_fname)

Sets save_folder and filename for saving final result and min_val result

Args:

save_folder (str | pathlib.Path): _description_ final_result_fname (str): _description_

Raises:

ValueError: _description_

Parameters:

• save_folder (str | Path)

• final_result_fname (str)

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

This function sets the checkpointing

Parameters:

• checkpoint_fname (str)

• checkpoint_freq (int)

• load_checkpoint_if_exists (bool)

• save_folder (str | Path)

check_training_ready(error=True, autofill=True, verbose=True)

This should always pass, all of these things are required to initialize a LIONsolver object

Args:

error (bool, optional): _description_. Defaults to True. autofill (bool, optional): _description_. Defaults to True.

Returns:

_type_: _description_

check_complete(error=True, autofill=True)

This function checks if the solver is complete, i.e. if all the necessary parameters are set to start traning.

save_checkpoint(epoch)

This function saves a checkpoint of the model and the optimizer

save_validation(epoch)

This function saves the validation results

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

This function saves the final results of the optimization

clean_checkpoints()

This function cleans the checkpoints

load_checkpoint()

This function loads a checkpoint (if exists)

train_step()

This function is responsible for performing a single tranining set epoch of the optimization. returns the average loss of the epoch

epoch_step(epoch)

This function is responsible for performing a single epoch of the optimization.

train(n_epochs)

This function is responsible for performing the optimization.

validate()

This function is responsible for performing a single validation set of the optimization. returns the average loss of the validation set this epoch.

get_model()

Returns the model used by the solver.

Return type:

LIONmodel

abstract mini_batch_step(sino_batch, target_batch)

This function should perform a single step of the optimization and return the loss

Return type:

Tensor