Matthieu Terris, Thomas Moreau, Nelly Pustelnik, Julián Tachella.
To appear in CVPR 2024
Enforcing equivariance of the denoiser to certain transformations within PnP/RED algorithms improves the stability and reconstruction quality of the algorithm.
We consider algorithms where gradients (or proximal operators) of explicit priors are replaced by denoisers; these algorithms typically take the form (in the case of PnP)
$$x_{k+1} = \text{D}(x_k - \gamma \nabla f(x_k)),$$where
$\text{D}$
$\mathcal{G}$
$\text{D}_{\mathcal{G}} = \frac{1}{|G|} \sum_{g \in \mathcal{G}} T_g^{-1} \text{D}(T_g)$
$$\begin{align*}
&\text{Sample } g_k \sim \mathcal{G} \\\
&\text{Set } \text{D}_{g_k}(x) = T_{g_k}^{-1} \text{D}(T_{g_k} x) \\\
&x_{k+1} = \text{D}_{g_k}(x_k - \gamma \nabla f(x_k)).
\end{align*}$$To reproduce the experiments, first download the test datasets and place them in your data folder. Next, update the config/config.json file to point to the correct data folder. There, there are two folders to specify:
ROOT_DATASET: the folder within which the CBSD68 and set3c datasets are located;PATH_MRI_DATA: the path to the fastMRI .pt dataset.
Then, you can run the following scripts to reproduce the experiments:
PnP (click to expand)
On the set3c dataset for the motion blur problem, with the drunet model:
python running_pnp.py --problem='motion_blur' --model_name='drunet' --rand_rotations=0 --dataset_name='set3c' --results_folder='table_4/' --compute_lip=0 --sigma_den=0.02 --noise_level=0.01
On the set3c dataset for the motion blur problem, with the drunet model, and with equivariance:
python running_pnp.py --problem='motion_blur' --model_name='drunet' --rand_rotations=1 --dataset_name='set3c' --results_folder='table_4/' --compute_lip=0 --sigma_den=0.02 --noise_level=0.01
RED (click to expand)
On the set3c dataset for the super-resolution blur problem, with the drunet model (Fig. 6 of the paper):
python running_red.py --problem='sr' --model_name='drunet' --rand_translations=0 --dataset_name='set3c' --sigma_den=0.015 --sr=2On the set3c dataset for the motion blur problem, with the drunet model, and with equivariance (Fig. 6 of the paper):
python running_red.py --problem='sr' --model_name='drunet' --rand_translations=1 --dataset_name='set3c' --sigma_den=0.015 --sr=2Feel free to change problem and models!
ULA (click to expand)
On the BSD68 dataset for the super-resolution blur problem, with the drunet model (Fig. 8 of the paper):
python running_ula.py --problem='motion_blur' --model_name='drunet' --rand_translations=0 --dataset_name='subset_BSD20' --sigma_den=0.019On the BSD10 dataset for the motion blur problem, with the drunet model, and with equivariance (Fig. 8 of the paper):
python running_ula.py --problem='motion_blur' --model_name='drunet' --rand_translations=1 --dataset_name='subset_BSD20' --sigma_den=0.019Feel free to change problem and models!
This code was tested with the following packages:
- torch 2.2
- deepinverse 0.1.1
The deepinverse package can be installed with pip install deepinverse or by cloning the repository.