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_device()

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

Selected GPU 0 with 18645.4375 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 0x7fdd6c8e33b0>

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]
  0%|          | 20/5000 [00:00<00:24, 199.23it/s]
  1%|          | 45/5000 [00:00<00:22, 225.22it/s]
  1%|▏         | 71/5000 [00:00<00:20, 237.18it/s]
  2%|▏         | 97/5000 [00:00<00:20, 243.05it/s]
  2%|▏         | 123/5000 [00:00<00:19, 246.37it/s]
  3%|▎         | 149/5000 [00:00<00:19, 248.34it/s]
  4%|▎         | 175/5000 [00:00<00:19, 249.25it/s]
  4%|▍         | 200/5000 [00:00<00:19, 249.32it/s]
  5%|▍         | 226/5000 [00:00<00:19, 250.07it/s]
  5%|▌         | 252/5000 [00:01<00:18, 250.60it/s]
  6%|▌         | 278/5000 [00:01<00:18, 251.00it/s]
  6%|▌         | 304/5000 [00:01<00:18, 251.21it/s]
  7%|▋         | 330/5000 [00:01<00:18, 251.32it/s]
  7%|▋         | 356/5000 [00:01<00:18, 251.54it/s]
  8%|▊         | 382/5000 [00:01<00:18, 251.61it/s]
  8%|▊         | 408/5000 [00:01<00:18, 251.62it/s]
  9%|▊         | 434/5000 [00:01<00:18, 251.80it/s]
  9%|▉         | 460/5000 [00:01<00:18, 251.23it/s]
 10%|▉         | 486/5000 [00:01<00:17, 251.48it/s]
 10%|█         | 512/5000 [00:02<00:17, 251.74it/s]
 11%|█         | 538/5000 [00:02<00:17, 251.86it/s]
 11%|█▏        | 564/5000 [00:02<00:17, 251.82it/s]
 12%|█▏        | 590/5000 [00:02<00:17, 252.03it/s]
 12%|█▏        | 616/5000 [00:02<00:17, 252.02it/s]
 13%|█▎        | 642/5000 [00:02<00:17, 252.10it/s]
 13%|█▎        | 668/5000 [00:02<00:17, 252.28it/s]
 14%|█▍        | 694/5000 [00:02<00:17, 250.18it/s]
 14%|█▍        | 720/5000 [00:02<00:17, 250.88it/s]
 15%|█▍        | 746/5000 [00:02<00:16, 251.50it/s]
 15%|█▌        | 772/5000 [00:03<00:16, 251.61it/s]
 16%|█▌        | 798/5000 [00:03<00:16, 251.84it/s]
 16%|█▋        | 824/5000 [00:03<00:16, 251.82it/s]
 17%|█▋        | 850/5000 [00:03<00:16, 251.68it/s]
 18%|█▊        | 876/5000 [00:03<00:16, 251.97it/s]
 18%|█▊        | 902/5000 [00:03<00:16, 252.01it/s]
 19%|█▊        | 928/5000 [00:03<00:16, 251.85it/s]
 19%|█▉        | 954/5000 [00:03<00:16, 251.90it/s]
 20%|█▉        | 980/5000 [00:03<00:15, 252.05it/s]
 20%|██        | 1006/5000 [00:04<00:15, 251.18it/s]
 21%|██        | 1032/5000 [00:04<00:15, 248.06it/s]
 21%|██        | 1057/5000 [00:04<00:16, 246.09it/s]
 22%|██▏       | 1082/5000 [00:04<00:16, 244.40it/s]
 22%|██▏       | 1107/5000 [00:04<00:16, 243.10it/s]
 23%|██▎       | 1132/5000 [00:04<00:15, 242.25it/s]
 23%|██▎       | 1157/5000 [00:04<00:15, 241.87it/s]
 24%|██▎       | 1182/5000 [00:04<00:15, 241.57it/s]
 24%|██▍       | 1207/5000 [00:04<00:15, 241.33it/s]
 25%|██▍       | 1232/5000 [00:04<00:15, 241.07it/s]
 25%|██▌       | 1257/5000 [00:05<00:15, 240.93it/s]
 26%|██▌       | 1282/5000 [00:05<00:15, 240.80it/s]
 26%|██▌       | 1307/5000 [00:05<00:15, 240.66it/s]
 27%|██▋       | 1332/5000 [00:05<00:15, 240.49it/s]
 27%|██▋       | 1357/5000 [00:05<00:15, 240.64it/s]
 28%|██▊       | 1382/5000 [00:05<00:15, 240.59it/s]
 28%|██▊       | 1407/5000 [00:05<00:14, 240.47it/s]
 29%|██▊       | 1432/5000 [00:05<00:14, 240.46it/s]
 29%|██▉       | 1457/5000 [00:05<00:14, 239.69it/s]
 30%|██▉       | 1482/5000 [00:06<00:14, 239.96it/s]
 30%|███       | 1507/5000 [00:06<00:14, 240.31it/s]
 31%|███       | 1532/5000 [00:06<00:14, 240.06it/s]
 31%|███       | 1557/5000 [00:06<00:14, 240.19it/s]
 32%|███▏      | 1582/5000 [00:06<00:14, 240.36it/s]
 32%|███▏      | 1607/5000 [00:06<00:14, 240.47it/s]
 33%|███▎      | 1632/5000 [00:06<00:14, 240.50it/s]
 33%|███▎      | 1657/5000 [00:06<00:13, 240.73it/s]
 34%|███▎      | 1682/5000 [00:06<00:13, 240.47it/s]
 34%|███▍      | 1707/5000 [00:06<00:13, 240.48it/s]
 35%|███▍      | 1732/5000 [00:07<00:13, 240.49it/s]
 35%|███▌      | 1757/5000 [00:07<00:13, 240.58it/s]
 36%|███▌      | 1782/5000 [00:07<00:13, 240.53it/s]
 36%|███▌      | 1807/5000 [00:07<00:13, 240.56it/s]
 37%|███▋      | 1832/5000 [00:07<00:13, 240.64it/s]
 37%|███▋      | 1857/5000 [00:07<00:13, 240.91it/s]
 38%|███▊      | 1882/5000 [00:07<00:12, 241.15it/s]
 38%|███▊      | 1907/5000 [00:07<00:12, 241.23it/s]
 39%|███▊      | 1932/5000 [00:07<00:12, 241.00it/s]
 39%|███▉      | 1957/5000 [00:07<00:12, 240.95it/s]
 40%|███▉      | 1982/5000 [00:08<00:12, 240.74it/s]
 40%|████      | 2007/5000 [00:08<00:12, 240.78it/s]
 41%|████      | 2032/5000 [00:08<00:12, 240.93it/s]
 41%|████      | 2057/5000 [00:08<00:12, 239.54it/s]
 42%|████▏     | 2082/5000 [00:08<00:12, 240.16it/s]
 42%|████▏     | 2107/5000 [00:08<00:12, 240.49it/s]
 43%|████▎     | 2132/5000 [00:08<00:11, 240.74it/s]
 43%|████▎     | 2157/5000 [00:08<00:11, 240.67it/s]
 44%|████▎     | 2182/5000 [00:08<00:11, 240.82it/s]
 44%|████▍     | 2207/5000 [00:09<00:11, 241.07it/s]
 45%|████▍     | 2232/5000 [00:09<00:11, 241.36it/s]
 45%|████▌     | 2257/5000 [00:09<00:11, 241.49it/s]
 46%|████▌     | 2282/5000 [00:09<00:11, 241.25it/s]
 46%|████▌     | 2307/5000 [00:09<00:11, 241.10it/s]
 47%|████▋     | 2332/5000 [00:09<00:11, 241.05it/s]
 47%|████▋     | 2357/5000 [00:09<00:10, 240.94it/s]
 48%|████▊     | 2382/5000 [00:09<00:10, 240.89it/s]
 48%|████▊     | 2407/5000 [00:09<00:10, 240.92it/s]
 49%|████▊     | 2432/5000 [00:09<00:10, 240.80it/s]
 49%|████▉     | 2457/5000 [00:10<00:10, 240.90it/s]
 50%|████▉     | 2482/5000 [00:10<00:10, 240.69it/s]
 50%|█████     | 2507/5000 [00:10<00:10, 240.81it/s]
 51%|█████     | 2532/5000 [00:10<00:10, 240.78it/s]
 51%|█████     | 2557/5000 [00:10<00:10, 240.84it/s]
 52%|█████▏    | 2582/5000 [00:10<00:10, 240.97it/s]
 52%|█████▏    | 2607/5000 [00:10<00:09, 240.90it/s]
 53%|█████▎    | 2632/5000 [00:10<00:09, 240.86it/s]
 53%|█████▎    | 2657/5000 [00:10<00:09, 240.80it/s]
 54%|█████▎    | 2682/5000 [00:10<00:09, 240.86it/s]
 54%|█████▍    | 2707/5000 [00:11<00:09, 240.74it/s]
 55%|█████▍    | 2732/5000 [00:11<00:09, 240.87it/s]
 55%|█████▌    | 2757/5000 [00:11<00:09, 241.01it/s]
 56%|█████▌    | 2782/5000 [00:11<00:09, 240.94it/s]
 56%|█████▌    | 2807/5000 [00:11<00:09, 241.08it/s]
 57%|█████▋    | 2832/5000 [00:11<00:08, 241.29it/s]
 57%|█████▋    | 2857/5000 [00:11<00:08, 241.18it/s]
 58%|█████▊    | 2882/5000 [00:11<00:08, 241.22it/s]
 58%|█████▊    | 2907/5000 [00:11<00:08, 241.18it/s]
 59%|█████▊    | 2932/5000 [00:12<00:08, 241.05it/s]
 59%|█████▉    | 2957/5000 [00:12<00:08, 241.14it/s]
 60%|█████▉    | 2982/5000 [00:12<00:08, 241.04it/s]
 60%|██████    | 3007/5000 [00:12<00:08, 241.14it/s]
 61%|██████    | 3032/5000 [00:12<00:08, 241.10it/s]
 61%|██████    | 3057/5000 [00:12<00:08, 240.99it/s]
 62%|██████▏   | 3082/5000 [00:12<00:07, 241.06it/s]
 62%|██████▏   | 3107/5000 [00:12<00:07, 241.01it/s]
 63%|██████▎   | 3132/5000 [00:12<00:07, 240.70it/s]
 63%|██████▎   | 3157/5000 [00:12<00:07, 240.69it/s]
 64%|██████▎   | 3182/5000 [00:13<00:07, 240.56it/s]
 64%|██████▍   | 3207/5000 [00:13<00:07, 240.56it/s]
 65%|██████▍   | 3232/5000 [00:13<00:07, 240.47it/s]
 65%|██████▌   | 3257/5000 [00:13<00:07, 240.43it/s]
 66%|██████▌   | 3282/5000 [00:13<00:07, 240.40it/s]
 66%|██████▌   | 3307/5000 [00:13<00:07, 240.49it/s]
 67%|██████▋   | 3332/5000 [00:13<00:06, 240.46it/s]
 67%|██████▋   | 3357/5000 [00:13<00:06, 240.41it/s]
 68%|██████▊   | 3382/5000 [00:13<00:06, 240.56it/s]
 68%|██████▊   | 3407/5000 [00:13<00:06, 240.48it/s]
 69%|██████▊   | 3432/5000 [00:14<00:06, 240.46it/s]
 69%|██████▉   | 3457/5000 [00:14<00:06, 240.58it/s]
 70%|██████▉   | 3482/5000 [00:14<00:06, 240.57it/s]
 70%|███████   | 3507/5000 [00:14<00:06, 240.54it/s]
 71%|███████   | 3532/5000 [00:14<00:06, 240.64it/s]
 71%|███████   | 3557/5000 [00:14<00:05, 240.70it/s]
 72%|███████▏  | 3582/5000 [00:14<00:05, 240.73it/s]
 72%|███████▏  | 3607/5000 [00:14<00:05, 240.79it/s]
 73%|███████▎  | 3632/5000 [00:14<00:05, 240.86it/s]
 73%|███████▎  | 3657/5000 [00:15<00:05, 240.93it/s]
 74%|███████▎  | 3682/5000 [00:15<00:05, 241.00it/s]
 74%|███████▍  | 3707/5000 [00:15<00:05, 241.00it/s]
 75%|███████▍  | 3732/5000 [00:15<00:05, 240.76it/s]
 75%|███████▌  | 3757/5000 [00:15<00:05, 240.91it/s]
 76%|███████▌  | 3782/5000 [00:15<00:05, 240.97it/s]
 76%|███████▌  | 3807/5000 [00:15<00:04, 240.97it/s]
 77%|███████▋  | 3832/5000 [00:15<00:04, 241.11it/s]
 77%|███████▋  | 3857/5000 [00:15<00:04, 241.12it/s]
 78%|███████▊  | 3882/5000 [00:15<00:04, 241.24it/s]
 78%|███████▊  | 3907/5000 [00:16<00:04, 240.95it/s]
 79%|███████▊  | 3932/5000 [00:16<00:04, 241.03it/s]
 79%|███████▉  | 3957/5000 [00:16<00:04, 241.04it/s]
 80%|███████▉  | 3982/5000 [00:16<00:04, 240.99it/s]
 80%|████████  | 4007/5000 [00:16<00:04, 241.07it/s]
 81%|████████  | 4032/5000 [00:16<00:04, 241.18it/s]
 81%|████████  | 4057/5000 [00:16<00:03, 241.21it/s]
 82%|████████▏ | 4082/5000 [00:16<00:03, 240.58it/s]
 82%|████████▏ | 4107/5000 [00:16<00:03, 240.70it/s]
 83%|████████▎ | 4132/5000 [00:17<00:03, 240.67it/s]
 83%|████████▎ | 4157/5000 [00:17<00:03, 240.35it/s]
 84%|████████▎ | 4182/5000 [00:17<00:03, 239.63it/s]
 84%|████████▍ | 4206/5000 [00:17<00:03, 239.50it/s]
 85%|████████▍ | 4230/5000 [00:17<00:03, 239.28it/s]
 85%|████████▌ | 4254/5000 [00:17<00:03, 239.33it/s]
 86%|████████▌ | 4278/5000 [00:17<00:03, 239.31it/s]
 86%|████████▌ | 4302/5000 [00:17<00:02, 239.34it/s]
 87%|████████▋ | 4326/5000 [00:17<00:02, 239.27it/s]
 87%|████████▋ | 4350/5000 [00:17<00:02, 239.24it/s]
 87%|████████▋ | 4374/5000 [00:18<00:02, 239.02it/s]
 88%|████████▊ | 4398/5000 [00:18<00:02, 239.04it/s]
 88%|████████▊ | 4422/5000 [00:18<00:02, 238.89it/s]
 89%|████████▉ | 4446/5000 [00:18<00:02, 238.73it/s]
 89%|████████▉ | 4470/5000 [00:18<00:02, 238.81it/s]
 90%|████████▉ | 4494/5000 [00:18<00:02, 238.81it/s]
 90%|█████████ | 4518/5000 [00:18<00:02, 238.15it/s]
 91%|█████████ | 4543/5000 [00:18<00:01, 238.93it/s]
 91%|█████████▏| 4567/5000 [00:18<00:01, 238.92it/s]
 92%|█████████▏| 4591/5000 [00:18<00:01, 239.24it/s]
 92%|█████████▏| 4615/5000 [00:19<00:01, 239.44it/s]
 93%|█████████▎| 4640/5000 [00:19<00:01, 239.71it/s]
 93%|█████████▎| 4665/5000 [00:19<00:01, 240.08it/s]
 94%|█████████▍| 4690/5000 [00:19<00:01, 240.03it/s]
 94%|█████████▍| 4715/5000 [00:19<00:01, 240.18it/s]
 95%|█████████▍| 4740/5000 [00:19<00:01, 240.49it/s]
 95%|█████████▌| 4765/5000 [00:19<00:00, 240.46it/s]
 96%|█████████▌| 4790/5000 [00:19<00:00, 240.70it/s]
 96%|█████████▋| 4815/5000 [00:19<00:00, 240.65it/s]
 97%|█████████▋| 4840/5000 [00:19<00:00, 240.65it/s]
 97%|█████████▋| 4865/5000 [00:20<00:00, 240.75it/s]
 98%|█████████▊| 4890/5000 [00:20<00:00, 240.68it/s]
 98%|█████████▊| 4915/5000 [00:20<00:00, 240.75it/s]
 99%|█████████▉| 4940/5000 [00:20<00:00, 240.81it/s]
 99%|█████████▉| 4965/5000 [00:20<00:00, 240.70it/s]
100%|█████████▉| 4990/5000 [00:20<00:00, 240.83it/s]
100%|██████████| 5000/5000 [00:20<00:00, 242.40it/s]
Iteration 4999, current converge crit. = 1.42E-05, objective = 1.00E-03
Iteration 4999, current converge crit. = 3.43E-04, objective = 1.00E-03
Linear reconstruction PSNR: 8.55 dB
Posterior mean PSNR: 22.32 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.