Uncertainty quantification with PnP-ULA.

This code shows you how to use sampling algorithms to quantify uncertainty of a reconstruction from incomplete and noisy measurements.

ULA obtains samples by running the following iteration:

\[x_{k+1} = x_k + \alpha \eta \nabla \log p_{\sigma}(x_k) + \eta \nabla \log p(y|x_k) + \sqrt{2 \eta} z_k\]

where \(z_k \sim \mathcal{N}(0, I)\) is a Gaussian random variable, \(\eta\) is the step size and \(\alpha\) is a parameter controlling the regularization.

The PnP-ULA method is described in the paper Laumont et al.[1].

import deepinv as dinv
from deepinv.utils.plotting import plot
import torch
from deepinv.utils import load_example

Load image from the internet

This example uses an image of Messi.

device = dinv.utils.get_freer_gpu() if torch.cuda.is_available() else "cpu"

x = load_example("messi.jpg", img_size=32).to(device)

Selected GPU 0 with 1951.125 MiB free memory

Define forward operator and noise model

This example uses inpainting as the forward operator and Gaussian noise as the noise model.

sigma = 0.1  # noise level
physics = dinv.physics.Inpainting(mask=0.5, img_size=x.shape[1:], device=device)
physics.noise_model = dinv.physics.GaussianNoise(sigma=sigma)

# Set the global random seed from pytorch to ensure reproducibility of the example.
torch.manual_seed(0)

<torch._C.Generator object at 0x7f3a5410bbd0>

Define the likelihood

Since the noise model is Gaussian, the negative log-likelihood is the L2 loss.

\[-\log p(y|x) \propto \frac{1}{2\sigma^2} \|y-Ax\|^2\]

# load Gaussian Likelihood
likelihood = dinv.optim.data_fidelity.L2(sigma=sigma)

Define the prior

The score a distribution can be approximated using Tweedie’s formula via the deepinv.optim.ScorePrior class.

\[\nabla \log p_{\sigma}(x) \approx \frac{1}{\sigma^2} \left(D(x,\sigma)-x\right)\]

This example uses a pretrained DnCNN model. From a Bayesian point of view, the score plays the role of the gradient of the negative log prior The hyperparameter sigma_denoiser (\(sigma\)) controls the strength of the prior.

In this example, we use a pretrained DnCNN model using the deepinv.loss.FNEJacobianSpectralNorm loss, which makes sure that the denoiser is firmly non-expansive (see Terris et al.[2]), and helps to stabilize the sampling algorithm.

sigma_denoiser = 2 / 255
prior = dinv.optim.ScorePrior(
    denoiser=dinv.models.DnCNN(pretrained="download_lipschitz")
).to(device)

Create the MCMC sampler

Here we use the Unadjusted Langevin Algorithm (ULA) to sample from the posterior defined in deepinv.sampling.ULAIterator. The hyperparameter step_size controls the step size of the MCMC sampler, regularization controls the strength of the prior and iterations controls the number of iterations of the sampler.

regularization = 0.9
step_size = 0.01 * (sigma**2)
iterations = int(5e3) if torch.cuda.is_available() else 10
params = {
    "step_size": step_size,
    "alpha": regularization,
    "sigma": sigma_denoiser,
}
f = dinv.sampling.sampling_builder(
    "ULA",
    prior=prior,
    data_fidelity=likelihood,
    max_iter=iterations,
    params_algo=params,
    thinning=1,
    verbose=True,
)

Generate the measurement

We apply the forward model to generate the noisy measurement.

y = physics(x)

Run sampling algorithm and plot results

The sampling algorithm returns the posterior mean and variance. We compare the posterior mean with a simple linear reconstruction.

mean, var = f.sample(y, physics)

# compute linear inverse
x_lin = physics.A_adjoint(y)

# compute PSNR
print(f"Linear reconstruction PSNR: {dinv.metric.PSNR()(x, x_lin).item():.2f} dB")
print(f"Posterior mean PSNR: {dinv.metric.PSNR()(x, mean).item():.2f} dB")

# plot results
error = (mean - x).abs().sum(dim=1).unsqueeze(1)  # per pixel average abs. error
std = var.sum(dim=1).unsqueeze(1).sqrt()  # per pixel average standard dev.
imgs = [x_lin, x, mean, std / std.flatten().max(), error / error.flatten().max()]
plot(
    imgs,
    titles=["measurement", "ground truth", "post. mean", "post. std", "abs. error"],
)

  0%|          | 0/5000 [00:00<?, ?it/s]
  1%|▏         | 66/5000 [00:00<00:07, 652.56it/s]
  3%|▎         | 133/5000 [00:00<00:07, 662.65it/s]
  4%|▍         | 200/5000 [00:00<00:07, 665.03it/s]
  5%|▌         | 267/5000 [00:00<00:07, 666.82it/s]
  7%|▋         | 334/5000 [00:00<00:06, 667.95it/s]
  8%|▊         | 405/5000 [00:00<00:06, 680.15it/s]
 10%|█         | 509/5000 [00:00<00:05, 796.13it/s]
 12%|█▏        | 611/5000 [00:00<00:05, 867.09it/s]
 14%|█▍        | 716/5000 [00:00<00:04, 922.13it/s]
 16%|█▋        | 821/5000 [00:01<00:04, 960.13it/s]
 19%|█▊        | 926/5000 [00:01<00:04, 985.77it/s]
 21%|██        | 1030/5000 [00:01<00:03, 999.64it/s]
 23%|██▎       | 1130/5000 [00:01<00:03, 996.10it/s]
 25%|██▍       | 1230/5000 [00:01<00:03, 988.01it/s]
 27%|██▋       | 1330/5000 [00:01<00:03, 989.43it/s]
 29%|██▊       | 1429/5000 [00:01<00:03, 984.21it/s]
 31%|███       | 1528/5000 [00:01<00:03, 981.18it/s]
 33%|███▎      | 1627/5000 [00:01<00:03, 983.24it/s]
 35%|███▍      | 1727/5000 [00:01<00:03, 985.51it/s]
 37%|███▋      | 1827/5000 [00:02<00:03, 987.50it/s]
 39%|███▊      | 1926/5000 [00:02<00:03, 986.54it/s]
 40%|████      | 2025/5000 [00:02<00:03, 987.23it/s]
 42%|████▎     | 2125/5000 [00:02<00:02, 988.36it/s]
 44%|████▍     | 2225/5000 [00:02<00:02, 989.35it/s]
 46%|████▋     | 2325/5000 [00:02<00:02, 989.60it/s]
 48%|████▊     | 2425/5000 [00:02<00:02, 990.45it/s]
 50%|█████     | 2525/5000 [00:02<00:02, 990.61it/s]
 52%|█████▎    | 2625/5000 [00:02<00:02, 990.82it/s]
 55%|█████▍    | 2725/5000 [00:02<00:02, 991.08it/s]
 56%|█████▋    | 2825/5000 [00:03<00:02, 990.99it/s]
 58%|█████▊    | 2925/5000 [00:03<00:02, 991.52it/s]
 60%|██████    | 3025/5000 [00:03<00:01, 991.52it/s]
 62%|██████▎   | 3125/5000 [00:03<00:01, 991.73it/s]
 64%|██████▍   | 3225/5000 [00:03<00:01, 985.44it/s]
 66%|██████▋   | 3325/5000 [00:03<00:01, 987.27it/s]
 68%|██████▊   | 3424/5000 [00:03<00:01, 983.78it/s]
 70%|███████   | 3523/5000 [00:03<00:01, 980.19it/s]
 72%|███████▏  | 3622/5000 [00:03<00:01, 980.50it/s]
 74%|███████▍  | 3721/5000 [00:03<00:01, 966.72it/s]
 76%|███████▋  | 3821/5000 [00:04<00:01, 973.78it/s]
 78%|███████▊  | 3921/5000 [00:04<00:01, 978.70it/s]
 80%|████████  | 4021/5000 [00:04<00:00, 983.12it/s]
 82%|████████▏ | 4121/5000 [00:04<00:00, 985.81it/s]
 84%|████████▍ | 4221/5000 [00:04<00:00, 988.02it/s]
 86%|████████▋ | 4321/5000 [00:04<00:00, 988.92it/s]
 88%|████████▊ | 4421/5000 [00:04<00:00, 989.53it/s]
 90%|█████████ | 4521/5000 [00:04<00:00, 989.71it/s]
 92%|█████████▏| 4621/5000 [00:04<00:00, 990.01it/s]
 94%|█████████▍| 4721/5000 [00:04<00:00, 990.15it/s]
 96%|█████████▋| 4821/5000 [00:05<00:00, 990.01it/s]
 98%|█████████▊| 4921/5000 [00:05<00:00, 990.48it/s]
100%|██████████| 5000/5000 [00:05<00:00, 955.42it/s]
Iteration 4999, current converge crit. = 1.43E-05, objective = 1.00E-03
Iteration 4999, current converge crit. = 3.42E-04, objective = 1.00E-03
Linear reconstruction PSNR: 8.55 dB
Posterior mean PSNR: 22.31 dB

References:

[1]

Rémi Laumont, Valentin De Bortoli, Andrés Almansa, Julie Delon, Alain Durmus, and Marcelo Pereyra. Bayesian imaging using plug & play priors: when langevin meets tweedie. SIAM Journal on Imaging Sciences, 15(2):701–737, 2022.

[2]

Matthieu Terris, Audrey Repetti, Jean-Christophe Pesquet, and Yves Wiaux. Building firmly nonexpansive convolutional neural networks. In ICASSP 2020-2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), 8658–8662. IEEE, 2020.