Supervised training

A solver trains a model: it is a torch training loop with what research needs around it (validation, checkpoints, saving with records, testing). This tutorial trains FBPConvNet on pairs of noisy sinograms and images with SupervisedSolver: first the least it needs, then each option.

Solvers are optional: a LION model is a torch module, which you can train with your own loop (the end of this tutorial). What LION’s solvers add is that everything about the training is recorded with the model, so it can be repeated, checked and published (the tutorials on records and publishing).

Short runs (a few epochs on a few slices) to show how; a real training needs many more.

Previous: Models. Next: Self-supervised training.

Needs a GPU and the processed LIDC-IDRI data in LION_DATA_PATH.

Note

The models trained in these tutorials are toys: a few epochs on a few slices, to show how LION is used, so their results are poor (often worse than FDK). Trained as their papers train them (many epochs, all the data), they do what the papers report.

The pieces

The data, the model, a torch optimizer and a loss. Any torch optimizer and loss work.

import tempfile
from pathlib import Path

import matplotlib.pyplot as plt
import torch
from torch.utils.data import DataLoader, Subset

from LION.experiments import ct_experiments
from LION.metrics.psnr import PSNR
from LION.models.post_processing.FBPConvNet import FBPConvNet
from LION.optimizers.SupervisedSolver import SupervisedSolver

experiment = ct_experiments.LowDoseCTRecon()
training = Subset(experiment.get_dataset("train"), range(40))
validation = Subset(experiment.get_dataset("validation"), range(8))
test = Subset(experiment.get_dataset("test"), range(8))

# torch DataLoaders: they batch (and shuffle) the samples
training_loader = DataLoader(training, batch_size=4, shuffle=True)
validation_loader = DataLoader(validation, batch_size=4)
test_loader = DataLoader(test, batch_size=4)

/mhome/damtp/r/ab2860/AItomotools/LION/data_loaders/LIDC_IDRI.py:487: UserWarning: LIDC-IDRI: left out 2 damaged or unusable files (re-run the pre-processing for these patients): ['LIDC-IDRI-0238 (not resampled)', 'LIDC-IDRI-0585 (not resampled)']
  warnings.warn(

The least a solver needs

The model, the optimizer and the loss; the training data; where to save the model. Then train(n_epochs) and save_final_results().

model = FBPConvNet(experiment.operator_params)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

solver = SupervisedSolver(model, optimizer, torch.nn.MSELoss())
solver.set_training(training_loader)
folder = Path(tempfile.mkdtemp())
solver.set_saving(folder, "FBPConvNet")

solver.train(2)
solver.save_final_results()  # FBPConvNet.pt (weights) and FBPConvNet.json
print("training loss per epoch:", solver.train_loss)

training loss per epoch: [0.22622674 0.03689291]

That is a complete training: the saved model holds its weights, its parameters and the record of how it was trained. Everything below is optional.

Options of the solver

Every solver takes the same arguments, in this order: the model, the optimizer, the loss, then, all optional: the operator’s parameters (taken from the model when it has them), the solver’s own parameters, a model regularisation (a penalty on the model added to the loss, e.g. for PnP denoisers: Denoisers and Plug-and-Play), a learning-rate scheduler, how much it prints, the device and the seed.

model = FBPConvNet(experiment.operator_params)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=2, gamma=0.5)

solver = SupervisedSolver(
    model,
    optimizer,
    torch.nn.MSELoss(),
    scheduler=scheduler,  # optional; steps every epoch (scheduler_step="batch": every batch)
    verbose=2,  # 0: nothing; 1 (default): one progress bar; 2: also each epoch's, and messages
    seed=0,  # the default: the same seed, the same model
)
solver.set_training(training_loader)
folder = Path(tempfile.mkdtemp())
solver.set_saving(folder, "FBPConvNet")

Supervised solver training on device: cuda:0

Validation (optional) runs every validation_freq epochs and keeps the best model: by the lowest value, or the highest for a score like the PSNR.

solver.set_validation(
    validation_loader,
    validation_freq=1,
    validation_fn=PSNR(),
    validation_fname="FBPConvNet_best",
)

Testing (optional) is a score on test data, run by solver.test(). The Evaluation and comparison evaluates properly, against other methods.

solver.set_testing(test_loader, PSNR())

Checkpoints

A checkpoint is everything needed to carry on training as if never stopped: the model’s weights, the optimizer’s state (e.g. Adam’s moments), the scheduler’s, the epoch, the losses so far and the random state. Training takes long, and jobs get stopped: with checkpoints, a stopped training loses at most checkpoint_freq epochs.

• When they are saved: every checkpoint_freq epochs, as checkpoint_fname with its * replaced by the epoch, in the save folder.

• When they are loaded: at the start of train(), the latest checkpoint matching the name is loaded, if there is one (and load_checkpoint_if_exists, True by default), and training goes on from its epoch. train(n) trains until n epochs are done in all, so a run started again with the same set-up finishes where it would have.

• By default, with set_saving alone, the solver checkpoints every 10 epochs as <name>_checkpoint_*.pt. set_checkpointing sets them:

solver.set_checkpointing("FBPConvNet_check_*.pt", checkpoint_freq=2)
# solver.set_checkpointing(..., load_checkpoint_if_exists=False)  # always start anew

Training

solver.train(4)
solver.save_final_results()

print("training loss per epoch:", solver.train_loss)
print("validation PSNR:", solver.validation_loss)
print("learning rate:", scheduler.get_last_lr())
print(sorted(p.name for p in folder.iterdir()))

Loading checkpoint...
No checkpoint matching /tmp/tmpttn3aguj/FBPConvNet_check_*.pt yet; starting from epoch 0.
Training seed: 0 (the solver's seed argument; seed=None for an unseeded run)
Validation loss: -19.538342475891113 - Validation loss std: 0.8199764012217944
Epoch 1 - Training loss: 0.19567729365080594 - Validation loss: -19.538342475891113
Validation loss: -22.70351815223694 - Validation loss std: 2.6133758422053686
Epoch 2 - Training loss: 0.03723795078694821 - Validation loss: -22.70351815223694
Validation loss: -0.12658442556858063 - Validation loss std: 2.0804106749230904
Epoch 3 - Training loss: 0.022894963435828686 - Validation loss: -0.12658442556858063
Validation loss: 26.368274927139282 - Validation loss std: 0.7803903163904746
Epoch 4 - Training loss: 0.01861754758283496 - Validation loss: 26.368274927139282
training loss per epoch: [0.19567729 0.03723795 0.02289496 0.01861755]
validation PSNR: [-19.53834248 -22.70351815  -0.12658443  26.36827493]
learning rate: [0.00025]
['FBPConvNet.json', 'FBPConvNet.pt', 'FBPConvNet_best.json', 'FBPConvNet_best.pt', 'FBPConvNet_check_0002.json', 'FBPConvNet_check_0002.pt', 'FBPConvNet_check_0004.json', 'FBPConvNet_check_0004.pt']

Resuming: a new solver with the same set-up (as after a crash) finds the checkpoint of epoch 4 and only trains epochs 5 and 6, with the losses of the first four kept.

model = FBPConvNet(
    experiment.operator_params
)  # new weights: replaced by the checkpoint's
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
resumed = SupervisedSolver(
    model,
    optimizer,
    torch.nn.MSELoss(),
    scheduler=torch.optim.lr_scheduler.StepLR(optimizer, step_size=2, gamma=0.5),
)
resumed.set_training(training_loader)
resumed.set_saving(folder, "FBPConvNet")
resumed.set_checkpointing("FBPConvNet_check_*.pt", checkpoint_freq=2)
resumed.train(6)
print("training loss per epoch:", resumed.train_loss)

training loss per epoch: [0.19567729 0.03723795 0.02289496 0.01861755 0.02060749 0.01653967]

Once a run is finished, its checkpoints can go:

resumed.clean_checkpoints()

Testing and what was saved

scores = solver.test()
print(f"test PSNR: {scores.mean():.2f} dB")

Testing loss: 27.081647396087646 - Testing loss std: 2.676541561270254
test PSNR: 27.08 dB

The final model and the best model by validation, each as a .pt (weights) and a .json (parameters and records). load makes the model again; the record says how it was trained (the data, the solver and its settings, the optimizer, the scheduler, the seed, the losses), enough to train the same model again (Saving, loading and records):

trained, options, data = FBPConvNet.load(folder / "FBPConvNet")
best, _, _ = FBPConvNet.load(folder / "FBPConvNet_best")
print(options.training.solver["_target_"], options.training.optimizer["lr"])

fig, ax = plt.subplots()
ax.plot(resumed.train_loss, label="training loss (MSE)")
ax.set_xlabel("epoch")
ax.legend(loc="upper left")
ax2 = ax.twinx()
ax2.plot(solver.validation_loss, "C1", label="validation PSNR")
ax2.legend(loc="upper right")

LION.optimizers.SupervisedSolver.SupervisedSolver 0.001

<matplotlib.legend.Legend object at 0x7f2a4a87b4d0>

Reconstructing

A reconstructor (LION.reconstructors) runs a trained model the way its solver did (the initial reconstruction, the normalisation), on one sample or a batch, without gradients. It is optional too: the model can be called directly, as below.

from LION.reconstructors.ForwardReconstructor import ForwardReconstructor

best.eval()  # the layers' evaluation behaviour (dropout, batch normalisation)
reconstructor = ForwardReconstructor(experiment.operator_params, best)
sinogram, image = validation[0]
reconstruction = reconstructor.reconstruct(sinogram)
# the same, by hand: FBPConvNet takes the measurements
# with torch.no_grad():
#     reconstruction = best(sinogram[None])[0]

fig, axes = plt.subplots(1, 2, figsize=(8, 4))
for ax, (title, x) in zip(
    axes, [("target", image), ("FBPConvNet, 4 epochs", reconstruction)]
):
    ax.imshow(x[0].cpu(), cmap="gray")
    ax.set_title(title)
    ax.axis("off")

Other supervised models

Any model is trained the same way: only the model changes. E.g. LPD:

from LION.models.iterative_unrolled.LPD import LPD
model = LPD(experiment.operator_params)

Without LION’s solver

A LION model is a torch module: train it with your own loop if you prefer. What you lose is what the solver records (which data, how it was trained), and so what the later tutorials do with it: training again from the record, checking that comparisons are fair, publishing.

from LION.utils.device import default_device

# the solver moved the model to LION's device (the GPU); by hand, you do it.
# The samples are already there.
model = FBPConvNet(experiment.operator_params).to(default_device())
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
model.train()
for sinograms, images in training_loader:
    optimizer.zero_grad()
    loss = torch.nn.functional.mse_loss(model(sinograms), images)
    loss.backward()
    optimizer.step()
print(f"one epoch by hand, last loss {loss.item():.4f}")

one epoch by hand, last loss 0.0232