This repository contains the implementations for the chapter Learned Regularization Functionals for Inverse Problems: A Comparative Study. If you have any questions or remarks, feel free to open an issue.
This readme file is structured into
- Installation instructions of the dependencies
- An overview of implemented regularizer architectures and instructions how to use them for reconstruction
- An overview of implemented training methods
- Instructions to reproduce the evaluation runs from the chapter
- Instructions to reproduce the training runs from the chapter
- Instructions to reproduce the automatic parameter fitting routine (Experiment 2) from the chapter
- A
bibtexcitation for the chapter
To get started using conda, clone the repository and run (may take a few minutes)
conda env create --file=environment.yaml
In particular, the code relies on the DeepInverse library, which can be installed as follows:
pip install deepinv
The implementation contains several regularizers (see table below) for which we provide trained weights. Creating a regularizer object consists out of three steps:
- Create the regularizer itself with the import and constructor stated in the first table below.
- For all architectures despite PatchNR, EPLL and LPN: use
from priors import ParameterLearningWrapperandregularizer = ParameterLearningWrapper(regularizer)to incorporate learned regularization parameters - Load the weights with
regularizer.load_state_dict(torch.load(path))with path stated in the paragraph "Weight Paths" below.
Then every regularizer has the fields regularizer.g(x) to evaluate the regularizer at x and regularizer.grad(x) to compute the gradient at x, where x is in the standard image format (batch size x channels x height x width).
Once the regularizer is loaded it can be used to solve an inverse problem using the nonmontonic accelerated (proximal) gradient descent (nmAPG) with the following code, where physics and data_fidelity provide the forward operator and data fidelity term in the format of the DeepInverse library:
from evaluation import reconstruct_nmAPG
# x_gt is the ground truth
y = physics(x_gt) # y is the observation
lmbd = 1 # regularization parameter
step_size = 0.1 # initial step size in the nmAPG
maxiter = 1000 # maximal number of iterations in the nmAPG
tol = 1e-4 # stopping criterion
recon = reconstruct_nmAPG(y, physics, data_fidelity, regularizer, lmbd, step_size, maxiter, tol)
| Regularizer Name | Import | Constructor |
|---|---|---|
| CRR | from priors import WCRR |
WCRR(sigma=0.1, weak_convexity=0.0) |
| WCRR | from priors import WCRR |
WCRR(sigma=0.1, weak_convexity=1.0) |
| ICNN | from priors import ICNNPrior |
ICNNPrior(in_channels=1, channels=32, kernel_size=5) |
| IDCNN | from priors import IDCNNPrior |
IDCNNPrior(in_channels=1, channels=32, kernel_size=5, act_name=act_name) where act_name = "elu" for the CT-AR example and act_name = "smoothed_relu" otherwise |
| CNN (for AR/LAR) | from priors import LocalAR |
LocalAR(in_channels=1, pad=True, use_bias=True, n_patches=-1, output_factor=output_factor) with output_factor= n_pixels / 321**2 where n_pixels is the number of pixels in the image to which the regularizer is applied |
| CNN (for bilevel) | from priors import LocalAR |
LocalAR(in_channels=1, pad=True, use_bias=False, n_patches=-1, reduction="sum", output_factor=1 / 142 ** 2) |
| TDV | from priors import TDV |
TDV(in_channels=1, num_features=32, multiplier=1, num_mb=3, num_scales=3, zero_mean=True) |
| LSR (despite NETT) | from priors import LSR |
LSR(nc=[32, 64, 128, 256], pretrained_denoiser=False, alpha=1.0, sigma=3e-2) |
| LSR (for NETT) | from priors import NETT |
NETT(in_channels=1, out_channels=1, hidden_channels=64, padding_mode="zeros") |
| EPLL | from priors import EPLL |
see below |
| PatchNR | from priors import PatchNR |
see below |
For PatchNR and EPLL the weight files also contain information about the architecture. We refer to the script variational_reconstruction.py (line 146 to 178) for an example how to load them.
We provide the weights for all experiments done in the chapter. To load the weights use torch.load(path) with path defined below, where problem="Denoising" or problem="CT" and regularizer_name is a string which defines the regularizer as in the above table.
- bilevel-JFB:
path=f"weights/bilevel_{problem}/{regularizer_name}_JFB_for_{problem}.pt" - bilevel-IFT:
path=f"weights/bilevel_{problem}/{regularizer_name}_IFT_for_{problem}.pt" - MAID:
path=f"weights/bilevel_{problem}/{regularizer_name}_IFT-MAID_for_{problem}.pt" - AR/LAR:
path=f"weights/adversarial_{problem}/{regularizer_name}_for_{problem}_fitted.pt" - other: the PatchNR weights are in the directory
weights/patchnr, the weights for EPLL are top-level in theweightsdirectory. The LPN weights are in the directoriesweights/lpn*
We proivde generic scripts for the different training routines. These training routines include:
- Bilevel training with IFT/JFB: use
from training_methods import bilevel_training, where the arguments are defined in the top of the scripttraining_methods/bilevel_training.py - Bilevel training with MAID: use
from training_methods import bilevel_training_maid, where the arguments are defined in the top of the scripttraining_methods/bilevel_training_maid.py - (Local) Adversarial Regularization: use
from training_methods import ar_training, where the arguments are defined in the top of the scripttraining_methods/ar_training.py - Score matching: use
from training_methods import score_training, where the arguments are defined in the top of the scripttraining_methods/score_training.py - Other: Similar there are the training routines for NETT, LPN, EPLL and PatchNR which are tailored to specific regularizers. The corresponding training scripts are top level (using
training_methods/nett_training.pyandtraining_methods/lpn_training.py)
We describe how the learned regularizers can be evaluated (e.g. to solve the variational problem). The code for reproducing the baselines is located in the baselines directory.
The weights used for generating the numbers in the chapter are contained in the repository. Alternatively, they can be regenerated as described in part 5 of the readme.
The script variational_reconstruction.py provides a unified evaluation routine for most of the regularizers included. An example command is the following:
python variational_reconstruction.py --problem Denoising --evaluation_mode IFT --regularizer_name CRR
The arguments can be chosen as follows:
--problemis set toDenoisingfor experiment 1 and toCTfor experiment 3--evaluation_modeis set toIFTfor bilevel-IFT,JFBfor bilevel-JFB,IFT-MAIDfor MAID,Scorefor score matching andARfor (local) adverserial regualrization--regularizer_namedefines the regularizer architecture. Valid names areCRR,WCRR,ICNN,IDCNN,LAR(referring to the CNN),TDV,LSR,PatchNR,EPLLandNETT.
Other comments:
- for
EPLLandPatchNRuse--evaluation_mode Score(even though that's not quite accurate) - for
NETTuse both--evaluation_mode NETTand--regularizer_name NETT - If you want to save the first 10 reconstruction you can add the flag
--save_results True
Even though the LPN provably defines a regularizer, evaluating it (or its gradient) requires to solve a (convex) optimization problem. Therefore, we evaluate the LPN in a Plug-and-Play fashion. In the denoising case this is a one-step reconstruction, in the CT case this is based on the ADMM algorithm.
To evaluate experiment 1 run:
python eval_LPN.py --problem Denoising --dataset BSD
To evaluate experiment 3 run:
python eval_LPN.py --problem CT --dataset LoDoPaB
To reproduce the training runs, we have unified scripts to reproduce all bilevel methods, all adversarial regularization methods and custom routines for NETT, EPLL, PatchNR and LPN,
For reproducing the bilevel results, use the script training_bilevel.py, e.g., as follows
python training_bilevel.py --problem Denoising --hypergradient IFT --regularizer_name CRR
where the arguments can be chosen as
--problemcan be eitherDenoising(experiment 1) orCT(experiment 3)--hypergradientcan beIFT,JFBorIFT-MAID--regularizer_namecan beCRR,WCRR,ICNN,IDCNN,LAR(for the CNN column),TDVorLSR.
For reproducing the AR/LAR runs, use the script training_AR.py, e.g., as follows
python training_AR.py --problem Denoising --regularizer_name CRR
where the arguments can be chosen as
--problemcan be eitherDenoising(experiment 1) orCT(experiment 3)--regularizer_namecan beCRR,WCRR,ICNN,IDCNN,LAR(for the CNN column) orTDV.
For experiment 1 the EPLL, PatchNR, LPN and NETT can be trained by the commands
python training_EPLL.py --problem Denoising
python training_patchnr.py --problem Denoising
python training_LPN.py --dataset BSD
python training_nett.py --problem Denoising
For experiment 3, the commands are
python training_EPLL.py --problem CT
python training_patchnr.py --problem CT
python training_LPN.py --dataset LoDoPaB
python training_NETT.py --problem CT
To reproduce the experiment 2 (training for denoising on BSDS and evaluating for CT on LoDoPaB), use the script parameter_fitting_Denoising_to_CT.py, e.g., by
python parameter_fitting_Denoising_to_CT.py --evaluation_mode bilevel-IFT --regularizer_name CRR
with arguments
--evaluation_modecan bebilevel-IFT,bilevel-JFB,AR,ScoreorNETT. UseScorefor EPLL and PatchNR (even though it's not quite accurate)--regularizer_namecan beCRR,WCRR,ICNN,IDCNN,LAR(for the CNN column),PatchNR,EPLL,TDV,LSRorNETT
Again, the LPN requires a custom routine. As common for Plug-and-Play methods the noise level during training partially determines the regularization strength. Therefore, we retrain the regularizer on denoising with a smaller noise level and then evaluate it on CT by calling
python training_LPN.py --dataset BSD --noise_level 0.05
python eval_LPN.py --problem CT --dataset BSD
@article{LearnedRegularizers,
title={Learning Regularization Functionals for Inverse Problems: A Comparative Study},
author={Johannes Hertrich and Hok Shing Wong and Alexander Denker and Stanislas Ducotterd and Zhenghan Fang and Markus Haltmeier and Željko Kereta and Erich Kobler and Oscar Leong and Mohammad Sadegh Salehi and Carola-Bibiane Schönlieb and Johannes Schwab and Zakhar Shumaylov and Jeremias Sulam and German Shâma Wache and Martin Zach and Yasi Zhang and Matthias J. Ehrhardt and Sebastian Neumayer},
year={2025},
journal={arXiv preprint arXiv:2510.01755},
}