LION: Learned Iterative Optimization Networks

LION is a PyTorch library for learned reconstruction in tomography and other inverse problems, made for research by the Cambridge Image Analysis (CIA) group at the University of Cambridge. It stays close to plain torch: models are torch modules, operators are functions of torch tensors, solvers are torch training loops, and any piece can be your own. Around that, it keeps what research needs to be repeated and compared:

• experiments: the geometry, noise and data of a problem, the same for every method;
• records saved with every model: the data it saw, how it was trained, the software it ran on, so it can be trained again or checked;
• evaluation and comparison of methods on the same test data, with checks that the comparison is fair, and the tables and figures a paper needs;
• publishing a trained model as a package anyone can load, test or train again.

Documentation: https://cambridgecia.github.io/LION/ (tutorials, the methods in LION, the API reference).

LION is approaching its first release (v1.0); until then its API may change.

Install

git clone https://github.com/CambridgeCIA/LION.git
cd LION
conda env create --file=env_base.yml --name=lion   # or another name
conda activate lion
pip install .

With mamba, use mamba env create --file=env_base.yml --name=lion and mamba activate lion instead.

Everything compiled, torch included, comes from conda (env_base.yml), so all of it uses the same C++ runtime. Don't install torch with pip on top: a pip torch loads the system's C++ runtime, and conda's libraries (scipy, ...) then fail with CXXABI_... not found on machines with an older system. Create the environment fresh from env_base.yml rather than updating an old one.

• CUDA version: set by cuda-version in env_base.yml (12.9). Any CUDA 12 driver runs it (nvidia-smi shows the driver's CUDA version).
• No GPU on the machine where you create the environment (e.g. a login node): conda cannot see a CUDA driver and refuses pytorch-gpu. Tell it which CUDA to assume: CONDA_OVERRIDE_CUDA=12.9 conda env create --file=env_base.yml --name=lion.

To develop LION, install it editable with the development tools, and the pre-commit hooks (Black formatting):

pip install -e ".[dev]"
pre-commit install --hook-type pre-commit --hook-type post-merge

A first training

import torch
from torch.utils.data import DataLoader

from LION.experiments import ct_experiments
from LION.models.post_processing.FBPConvNet import FBPConvNet
from LION.optimizers.SupervisedSolver import SupervisedSolver

# the problem: low-dose CT of LIDC-IDRI slices (geometry, noise, data)
experiment = ct_experiments.LowDoseCTRecon()
training_loader = DataLoader(experiment.get_dataset("train"), batch_size=4, shuffle=True)

# a model, made from the problem's geometry, and a torch optimizer and loss
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)
solver.set_saving("results/fbpconvnet", "FBPConvNet")
solver.train(100)
solver.save_final_results()  # the weights, and the records to train it again

Then evaluate and compare it (LION.evaluation), reconstruct with it (LION.reconstructors) or publish it (LION.package): the tutorials go through each, and scripts/example_scripts has complete scripts, one feature each.

Data and results

LION finds datasets in LION_DATA_PATH; set it, e.g. in your .bashrc: export LION_DATA_PATH=/path/to/Datasets. The supported datasets are LIDC-IDRI and 2DeteCT: how to download and prepare them.

LION_RESULTS_PATH (optional) is where results go when no folder is given. Paths inside either folder are written relative to it in LION's saved files, so the files mean the same on any machine.

NOTE: In CIA at DAMTP, the datasets are already there: export LION_DATA_PATH=/store/cia-lion/datasets, and set LION_RESULTS_PATH to your own user folder.

Choosing a GPU

LION uses the GPU if there is one. Set another once, at the top of your script, and LION uses it wherever no device is given (solvers, loading models, datasets, noise simulation):

import torch
torch.cuda.set_device(2)            # or: torch.set_default_device("cuda:2")

On machines with several GPUs where only one may be used, the most robust option is to make only that GPU visible to the process, so nothing (LION, PyTorch or ASTRA) can use another: CUDA_VISIBLE_DEVICES=2 python my_script.py.

What is where

Folder	What it has
LION/operators/	Forward operators: ct/ (the CT geometry and projector, noise, FDK, SIRT, TV), pcm/ (photocurrent mapping), wavelets, ...
LION/data_loaders/	Datasets (LIDC-IDRI, 2DeteCT) and their records
LION/experiments/	Experiments: a problem's geometry, noise and data
LION/models/	Models, by kind: post-processing, unrolled, learned FBP, CNNs and denoisers, PnP, learned regularisers
LION/optimizers/	Solvers: supervised, self-supervised, denoiser and adversarial training
LION/reconstructors/	Reconstructing with trained models: forward, PnP, learned regulariser
LION/classical_algorithms/	Conjugate gradient, FISTA, SPGL1, for any operator
LION/evaluation/, LION/metrics/	Evaluation and comparison; PSNR, SSIM, HaarPSI
LION/losses/	Losses and penalties on models
LION/package.py	Publishing and loading trained models
demos/	The tutorials
scripts/	Example scripts, and scripts that reproduce papers
docs/	The documentation site

Which methods LION has, and their papers: Methods in LION. How models record their making, and how to make them again: Records.

Not only CT

LION's general code uses general words, as an operator can be anything. For CT, they mean:

LION says	In CT
operator_params (an OperatorParameters)	the geometry (Geometry)
the operator (Operator, operator_params.make_operator())	the projector (CTProjectionOp)
domain (op.domain_shape)	the image (geometry.image_shape)
range (op.range_shape), measurements	the sinogram
the operator's default reconstruction	FDK

So a CT model is made from its geometry, e.g. LPD(geometry), and a solver takes operator_params=geometry. More about operators.

Contributing

Read the developer guide for how LION is built and what a new model, solver or operator needs, and the contributor guide for the conventions the code follows. Run the tests with pytest pytests/; tests that need CUDA or tomosipo/ASTRA are skipped on machines without them, and the same tests run on every pull request.

Citing LION

A paper is to come; until then, please cite the repository:

@software{LION,
  title  = {{LION}: Learned Iterative Optimization Networks},
  author = {Biguri, Ander and Shoebridge, Charlie and Runkel, Christina and Vu, Thanh Trung and Valat, Emilien and Sherry, Ferdia and Tan, Hong Ye and Sainz Bear, Andrea and Kiss, Maximilian B.},
  url    = {https://github.com/CambridgeCIA/LION},
}

and the papers of the methods you use (Model.cite()). GitHub's "Cite this repository" button gives the same, from CITATION.cff.

Licence

LION is released under the GNU General Public License v3 (GPL-3), see LICENSE.txt. It has to be: its CT backend, ASTRA and tomosipo, are GPL-3.

Contributors

Biguri 💻 🎨 🤔 ✅	Emilien Valat 💻 🎨 🤔 🔣	Ferdia 🎨 🔧 💻	michi217 💻 🔣	Oliver Coughlan 💻	mbkiss 🎨 💻	ChristinaRunkel 💻 🐛 🎨 💡 🤔 🔬 📓
Zak Shumaylov 💻	Hong Ye Tan 💻 🎨	Charlie Shoebridge 💻 🎨 🤔	Andrea Sainz Bear 💻 🐛 🔬	Thanh Trung Vu 💻 🐛 🔬 🎨	cslht 💻