Plug and Play optimization¶

Introduction¶

Plug-and-play (PnP) is an optimization framework that integrates modern denoising priors into gradient or proximal descent schemes.

The goal of this notebook is to use PnP methods to solve image restoration problems (interpolation of missing pixels, denoising, or deblurring for instance). Typically, these inverse problems can be written $$y = Mx+\eta$$ where $x$ is an unknown image that we wish to reconstruct, $y$ is the observed image, $M$ is a linear degradation operator and $\eta$ is a realization of a known noise distribution (for instance Gaussian). The discrete images $x$ and $y$ can be seen as vectors of size $n$ and $m$ (for instance by reading them columnwise), and $M$ is a $m\times n$ matrix. For instance, for an inverse problem where pixel values are missing, $M$ is a diagonal matrix with $0$ (hidden pixels) and $1$ (known pixels) on the diagonal.

To estimate $x$, when the noise is Gaussian, a classical approach consists in minimizing an energy of the form $$ \frac{1}{2\sigma^2}\|Mx - y \|^2 + U(x)$$
where the second term is a regularization term on $x$. In a Bayesian framework, this energy can be interpreted as $$-\log p(x|y) = -\log p(y|x)-\log p(x),$$ where $p(y|x) \propto e^{-\frac{1}{2\sigma^2}\|Mx - y \|^2}$ is the likelihood of $y$ knowing $x$ and a prior distribution $p(x) \propto e^{-U(x)}$ is assumed on the unknown $x$.

Typicall optimization schemes used to minimize the previous expression involve the gradient or proximal operator of $U$. The principle of PnP approaches is to derive an approximation of $\nabla U$ or $prox_U$ using an image denoiser, for instance a denoising deep neural network. This approximation can then be used with any scheme using gradient or proximal descent on $U$ for optimization or sampling.

This practical session explains and shows inpainting et deblurring experiments based on the code of the following paper:

E. K. Ryu, J. Liu, S. Wang, X. Chen, Z. Wang, and W. Yin, Plug-and-Play Methods Provably Converge with Properly Trained Denoisers, in Proceedings of the 36th International Conference on Machine Learning, ICML 2019, 9-15 June 2019, Long Beach, California, USA, 2019, pp. 5546–5557

References:

• Model and code for loading it: https://github.com/uclaopt/Provable_Plug_and_Play
• Notations for the notebook are the one of the paper R.Laumont, V. de Bortoli, A. Almansa, J. Delon, A. Durmus, M. Pereyra, Bayesian imaging using Plug & Play priors: when Langevin meets Tweedie, 2022

Authors of the notebook:

• Lucía Bouza
• Julie Delon
• Rémi Laumont

Environment setting¶

Importing packages¶

Below is a list of packages needed. PyTorch version used to run this notebook is 1.11.0+cu113 (to check the installed version, use torch.__version__)

• numpy
• matplotlib.pyplot (display of images and graphics)
• torch (use cuda with PyTorch)
• time (measure time)
• os (interactions with the operating system)

To import the solutions, execute the following cell. If you are using a Windows system, comment the wget line, download the file by hand, and place it in the same folder as the notebook.

Set a device¶

Next, we need to choose which device to run the algorithm on. Running the algorithm on large images takes longer and will go much faster when running on a GPU. We can use torch.cuda.is_available() to detect if there is a GPU available. Next, we set the torch.device. The .to(device) method is used to move tensors or modules to a desired device, we will use it in next sections.

Download code and models¶

In this notebook, for the sake of simplicity (and computing time!), we will use a pre-trained denoising networks. The load_denoiser function helps to load the model available:

• DnCNN with real spectral normalization (RealSN-DnCNN), presented in this paper. The model was taken from the original implementation.

In AuxiliarFunctions_PnP we define the network, taken from the original code.

Performance measurement¶

In the following, we will compare several algorithms, in terms of computing time and performance. We will use the PSNR (defined below) to measure restoration performance. The higher the value, the better the result.

Helper functions¶

In the next cell we will define some helper functions for displaying and creating degraded images. We will also define the Denoiser function to be used by all PnP implementations.

Algorithms¶

As explained above, we wish to solve $$ \arg\min_x f(x) + \alpha U(x)$$

where,

• $x$ is an image
• $f(x)$ is a data fidelity term (corresponding to the logarithm of the likelihood)
• $U(x)$ is the image prior
• $\alpha ≥ 0$ is the relative importance between $f$ and $g$

For convenience we can rewrite the minimization problem as: $$\arg\min_x \frac{1}{\alpha} f(x) + U(x).$$

PnP ADMM¶

Recall the alternating directions method of multipliers (ADMM) optimization to solve the previous problem:

$$ x^{k+1} =Prox_{\frac \epsilon \alpha f}(y^k −u^k) \\ y^{k+1} = Prox_{\epsilon U} (x^{k+1} + u^k ) \\ u^{k+1} = u^k + x^{k+1} − y^{k+1} \\ $$

PnP-ADMM simply replaces the proximal operator $Prox_{\epsilon U}$ with the denoiser $D_{\epsilon}$, where $\epsilon$ is the noise parameter (variance) of the denoiser:

$$ x^{k+1} =Prox_{\frac \epsilon \alpha f}(y^k −u^k) \\ y^{k+1} = D_{\epsilon} (x^{k+1} + u^k ) \\ u^{k+1} = u^k + x^{k+1} − y^{k+1}. \\ $$

Write a function pnp_admm to implement the algorithm, with the following parameters:

def pnp_admm(noisy, denoiser, proximal_step, **opts):
    """
    Parameters:
        :noisy - the noisy observation.
        :denoiser - the Gaussian denoiser used in Plug-and-Play ADMM.
        :proximal_step - the function which implements the proximal step of the ADMM algorithm.
        :opts - the kwargs for hyperparameters in Plug-and-Play ADMM.
    """

PnP FBS¶

The forward-backward splitting (FBS) optimization scheme can be written:

$ x^{k+1} = Prox_{\epsilon U}(I - \frac \epsilon \alpha \nabla f)(x^k) $

PNP-FBS simply replaces the proximal operator $Prox_{\epsilon U}$ with the denoiser $D_{\epsilon}$:

$ x^{k+1} = D_{\epsilon}(I - \frac \epsilon \alpha \nabla f)(x^k) $

Write a function pnp_fbs to implement the algorithm, with the following parameters:

def pnp_fbs(noisy, denoiser, proximal_step, **opts):
    """
    Parameters:
        :noisy - the noisy observation.
        :denoiser - the Gaussian denoiser used in Plug-and-Play FBS.
        :gradient_step - the function which implements the gradient step: x- alpha*grad(f)
        :opts - the kwargs for hyperparameters in Plug-and-Play FBS.
    """

PnP BBS¶

A fully proximal version of the algorithm, called backward-backward splitting (BBS) writes:

$ x^{k+1} = Prox_{\epsilon U}(Prox_{\frac \epsilon \alpha f})(x^k) $

PNP-BBS replaces the proximal operator $Prox_{\epsilon U}$ with the denoiser $ D_{\epsilon}$:

$ x^{k+1} = D_{\epsilon}(Prox_{\frac \epsilon \alpha f})(x^k) $

Write a function pnp_fbs to implement the algorithm, with the following parameters:

def pnp_bbs(noisy, denoiser, proximal_step, **opts):
    """
    Parameters:
        :noisy - the noisy observation.
        :denoiser - the Gaussian denoiser used in Plug-and-Play BBS.
        :proximal_step - the function which implements the proximal step of the FBS algorithm.
        :opts - the kwargs for hyperparameters in Plug-and-Play BBS.
    """

Remark: it can be shown that PNP-FBS and PNP-ADMM are different methods for finding the same set of fixed points. Sometimes, PNP-FBS is easier to implement since it only requires the computation of $\nabla f$ rather than $Prox_{ f}$ . On the other hand, PNP-ADMM is generally more robust with better convergence properties. BBS aims at solving a slightly modified version of the original problem where $f$ is replaced by its Moreau envelope.

Experiments¶

We will now use the previous algorithms to solve two inverse problems in imaging: deblurring and inpainting.

PnP for deblurring¶

In this section, experiments will be carried out to recover images with blur and additive noise. In the next cell we will define the algorithm's hyperparameters, load the denoiser model to use, and create a blurred image. In all cases we will use the denoising model developed by the authors (RealSN_DnCNN) with sigma level of 5.

To use the previous algorithms, we need to define the gradient and proximal steps for the data term of the deblurring problem. The data term can be written $$\frac 1 \alpha f(x) = \frac{|| Mx - y ||^2}{2 \alpha\sigma^2}$$ with,

• $M$ the degradation matrix (corresponding to a convolution with a kernel $h$)
• $y$ the observation
• $\sigma$ the standard deviation of the Gaussian additive noise for the given restoration problem.

Clearly, $$\nabla f(x) = M^*\frac{Mx -y}{\sigma^2}$$ and $$\mathrm{prox}_{\frac \epsilon \alpha f}(x) = (Id+\frac \epsilon {\alpha\sigma^2} M^* M)^{-1}\left(x +\frac \epsilon {\alpha\sigma^2} M^* y\right).$$

The linear operator $M$ represents a convolution by a kernel $h$. It can be applied directly in the Fourier domain as a multiplication. The adjoint $M^*$ is the convolution with the conjugate of $h$.

PnP ADMM¶

In the following cell we will define the function that implements the proximal step needed by PnP ADMM. Next, we will run the algorithm and display the result.

PnP-FBS¶

PnP BBS¶

1. Run the different algorithms for different values of $\alpha$ and number of iterations. What do you observe?
2. Are the results of the three algorithms identical? Describe what you observe.
3. Display the graph of the PSNR values along the iterations for each algorithm and comment.
4. How can we deal with the case $\sigma = 0$?

PnP for missing pixels¶

Everything works fine in the deblurring experiments above and all algorithms converge in a few dozens of iterations. We will now deal with an inverse problem where pixels are missing, and we will see that the same algorithms require much more iterations to converge.

We create the degraded image.

The following cell contains the gradient and proximal operators for the inpainting problem.

1. Run the different algorithms for different values of $\alpha$ and number of iterations. What do you observe?
2. Are the results of the three algorithms identical? Describe what you observe.

3.Display the graph of the PSNR values along the iterations for each algorithm and comment. 4. How can we deal with the case $\sigma = 0$?

PnADMM¶

PnP FBS¶

PnP-BBS¶