#### NOTE
New to DeepInverse? Get started with the basics with the
[5 minute quickstart tutorial](https://deepinv.org/auto_examples/basics/demo_quickstart.html.md#sphx-glr-auto-examples-basics-demo-quickstart-py)..
# DPS – Posterior Sampling with Diffusion Models
In this tutorial, we will go over the steps in the Diffusion Posterior Sampling (DPS) algorithm introduced in
Chung *et al.*[1](#footcite-chung2022diffusion). The full algorithm is implemented in [`deepinv.sampling.DPS`](https://deepinv.org/api/stubs/deepinv.sampling.DPS.html.md#deepinv.sampling.DPS).
Let us import the relevant modules and load a sample
image of size 64 x 64. This will be used as our ground truth image.
#### NOTE
We work with an image of size 64 x 64 to reduce the computational time of this example.
```Python
import torch
import deepinv as dinv
from deepinv.utils.plotting import plot
from deepinv.utils import load_example
import matplotlib as mpl
mpl.rcParams["animation.html"] = "jshtml"
device = dinv.utils.get_device()
x_true = load_example("butterfly.png", img_size=64, device=device)
x = x_true.clone()
```
```none
Selected GPU 0 with 24087.25 MiB free memory
```
In this tutorial we consider random inpainting as the inverse problem, where the forward operator is implemented
in [`deepinv.physics.Inpainting`](https://deepinv.org/api/stubs/deepinv.physics.Inpainting.html.md#deepinv.physics.Inpainting). In the example that we use, 90% of the pixels will be masked out randomly,
and we will additionally have Additive White Gaussian Noise (AWGN) of standard deviation 12.75/255.
```Python
sigma = 12.75 / 255.0 # noise level
physics = dinv.physics.Inpainting(
img_size=(3, x.shape[-2], x.shape[-1]),
mask=0.1,
pixelwise=True,
noise_model=dinv.physics.GaussianNoise(sigma),
device=device,
)
y = physics(x_true)
plot(
{
"Measurement": y,
"Ground Truth": x_true,
}
)
```
## Load a pre-trained denoiser
Our DPS implementation relies on a pre-trained denoiser, which is used to approximate the score function of the diffusion process. In this example, we will use a DRUNet denoiser, which is a widely used architecture for image denoising. The example should work with any other denoiser, as long as it takes as input an image and a noise level, and outputs a denoised image.
```Python
denoiser = dinv.models.DRUNet(device=device)
```
## The diffusion schedule
Our DPS implementation supports two standard diffusion schedules, which are the [`deepinv.sampling.VariancePreservingDiffusion`](https://deepinv.org/api/stubs/deepinv.sampling.VariancePreservingDiffusion.html.md#deepinv.sampling.VariancePreservingDiffusion) (VP) and [`deepinv.sampling.VarianceExplodingDiffusion`](https://deepinv.org/api/stubs/deepinv.sampling.VarianceExplodingDiffusion.html.md#deepinv.sampling.VarianceExplodingDiffusion) (VE) SDEs. In this example, we will use the VP SDE, which is the continuous-time limit of the DDPM sampling process.
## DPS approximation
In order to perform gradient-based **posterior sampling** with diffusion models, we have to be able to compute
$\nabla_{\mathbf{x}_t} \log p(\mathbf{x}_t|\mathbf{y})$. Applying Bayes rule, we have
$$
\nabla_{\mathbf{x}_t} \log p(\mathbf{x}_t|\mathbf{y}) = \nabla_{\mathbf{x}_t} \log p(\mathbf{x}_t)
+ \nabla_{\mathbf{x}_t} \log p(\mathbf{y}|\mathbf{x}_t)
$$
For the former term, we can simply plug-in our estimated score function as in Tweedie’s formula. As the latter term
is intractable, DPS proposes the following approximation (for details, see Theorem 1 of Chung *et al.*[1](#footcite-chung2022diffusion))
$$
\nabla_{\mathbf{x}_t} \log p(\mathbf{x}_t|\mathbf{y}) \approx \nabla_{\mathbf{x}_t} \log p(\mathbf{x}_t)
+ \nabla_{\mathbf{x}_t} \log p(\mathbf{y}|\widehat{\mathbf{x}}_{0}(\mathbf{x_t}))
$$
where $\widehat{\mathbf{x}}_{0}(\mathbf{x_t})$ is the posterior mean of the clean image given the noisy image at time $t$, which can be estimated with a denoiser network.
Under the assumption of Gaussian noise, the likelihood term can be written as
$$
\log p(\mathbf{y}|\widehat{\mathbf{x}}_0(\mathbf{x_t})) =
-\frac{\|\mathbf{y} - A\widehat{\mathbf{x}}_0(\mathbf{x_t})\|_2^2}{2\sigma_y^2}.
$$
Taking the gradient w.r.t. $\mathbf{x}_t$ requires backpropagation through the denoiser, which can be easily implemented with PyTorch’s autograd.
We provide an implementation of this approximation in [`deepinv.sampling.DPSDataFidelity`](https://deepinv.org/api/stubs/deepinv.sampling.DPSDataFidelity.html.md#deepinv.sampling.DPSDataFidelity), which is a subclass of [`deepinv.sampling.NoisyDataFidelity`](https://deepinv.org/api/stubs/deepinv.sampling.NoisyDataFidelity.html.md#deepinv.sampling.NoisyDataFidelity).
#### NOTE
The DPS algorithm assumes that the images are in the range [-1, 1], whereas standard denoisers
usually output images in the range [0, 1]. This is why we rescale the images before applying the steps.
```Python
from deepinv.sampling import DPSDataFidelity
x0 = x_true * 2.0 - 1.0 # [0, 1] -> [-1, 1]
data_fidelity = DPSDataFidelity(denoiser=denoiser, clip=(-1.0, 1.0))
# choose some arbitrary noise level
sigma_t = 0.2
xt = x0 + sigma_t * torch.randn_like(x0)
# DPS
grad, x0_t = data_fidelity.grad(
xt / 2 + 0.5, y=y, physics=physics, sigma=sigma_t / 2, get_model_outputs=True
) # Set get_model_outputs to True to also retrieve the denoised output
# Visualize
plot(
{
"Ground Truth": x0,
"Noisy": xt,
"Posterior Mean": x0_t,
"Gradient": grad,
}
)
```
## DPS Algorithm
As we visited all the key components of DPS, we are now ready to define the algorithm. For every denoising
timestep, the algorithm iterates the following
1. Get $\hat{\mathbf{x}}$ using the denoiser network.
2. Compute $\nabla_{\mathbf{x}_t} \log p(\mathbf{y}|\hat{\mathbf{x}}_t)$ through backpropagation.
3. Perform reverse diffusion sampling with DDPM(IM), corresponding to an update with $\nabla_{\mathbf{x}_t} \log p(\mathbf{x}_t)$.
4. Take a gradient step with $\nabla_{\mathbf{x}_t} \log p(\mathbf{y}|\hat{\mathbf{x}}_t)$.
There are two caveats here. First, in the original work, DPS used DDPM ancestral sampling. The ancestral solver
[`deepinv.sampling.AncestralSolver`](https://deepinv.org/api/stubs/deepinv.sampling.AncestralSolver.html.md#deepinv.sampling.AncestralSolver) retrieves DDPM when $\alpha = 1.0$, and the deterministic DDIM sampler
Song *et al.*[2](#footcite-song2020denoising) when $\alpha = 0$.
One can freely choose the $\alpha$ parameter here,
it is advisable to keep it $\alpha = 1.0$ if `num_steps=1000`.
Second, one can also switch to other diffusion schedules, such as the VE SDE, which corresponds to a different noise schedule and sampling process. In this case, the DPS approximation still holds, but the sampling step will be different.
With DeepInverse, we can use the [`deepinv.sampling.DPS`](https://deepinv.org/api/stubs/deepinv.sampling.DPS.html.md#deepinv.sampling.DPS) class to perform the above steps with minimal code, with some important parameters:
> - `weight`: corresponds to the $\lambda$ parameter in the above equation, which controls the strength of the gradient step.
> - `alpha`: the weight of the noise in the reverse-time SDE, which controls the strength of the noise in the reverse diffusion sampling step (1 for DDPM, 0 for DDIM).
> - `num_steps`: corresponds to the number of denoising steps, which is usually set to 1000 for best performance, but can be reduced to 200 for faster sampling.
#### NOTE
For simplicity, we only show the DPS with the VP / VE SDEs, but the algorithm can be easily adapted to **arbitrary** diffusion processes,
for example [`deepinv.sampling.EDMDiffusionSDE`](https://deepinv.org/api/stubs/deepinv.sampling.EDMDiffusionSDE.html.md#deepinv.sampling.EDMDiffusionSDE) with custom noise schedules.
Please refer to the example [Building your diffusion posterior sampling method using SDEs](https://deepinv.org/auto_examples/sampling/demo_diffusion_sde.html.md#sphx-glr-auto-examples-sampling-demo-diffusion-sde-py) for a full demonstration of how to modify the
algorithm.
#### NOTE
We only use 200 steps to reduce the computational time of this example. As suggested by the authors of DPS, the
algorithm works best with `num_steps = 1000`.
```Python
# Instantiate the model
model = dinv.sampling.DPS(
denoiser=denoiser,
schedule="vp",
num_steps=200,
weight=20.0,
guidance="annealed",
alpha=0.5,
verbose=True,
device=device,
dtype=torch.float64,
rng=torch.Generator(device=device),
)
# Run the sampling
with torch.no_grad():
sample, trajectory = model(
y.clone(),
physics,
seed=123, # for reproducibility!
get_trajectory=True,
)
# plot the results
plot(
{
"Measurement": y,
"Model Output": sample,
"Ground Truth": x_true,
}
)
anim = dinv.utils.plot_videos(
trajectory[::10],
time_dim=0,
suptitle="DPS Trajectory",
return_anim=True,
)
anim
```
```none
0%| | 0/199 [00:00
* **References:**
* **[1]** Hyungjin Chung, Jeongsol Kim, Michael Thompson Mccann, Marc Louis Klasky, and Jong Chul Ye. Diffusion posterior sampling for general noisy inverse problems. In *The Eleventh International Conference on Learning Representations*. 2022.
* **[2]** Jiaming Song, Chenlin Meng, and Stefano Ermon. Denoising diffusion implicit models. In *International Conference on Learning Representations*. 2020.
**Total running time of the script:** (0 minutes 2.101 seconds)
🏷 Tags: Inpainting
[`Download Jupyter notebook: demo_dps.ipynb`](https://deepinv.org/auto_examples/sampling/_downloads/31f92417d030e0d2382c52178325ad67/demo_dps.ipynb)
[`Download Python source code: demo_dps.py`](https://deepinv.org/auto_examples/sampling/_downloads/faf1222544a1ffa01f52a6b80e689f71/demo_dps.py)
[`Download zipped: demo_dps.zip`](https://deepinv.org/auto_examples/sampling/_downloads/f85d9f98328661ab41ee319eaca58073/demo_dps.zip)
[Gallery generated by Sphinx-Gallery](https://sphinx-gallery.github.io)