<!-- DO NOT EDIT. -->
<!-- THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. -->
<!-- TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: -->
<!-- "auto_examples/unfolded/demo_unfolded_constrained_LISTA.py" -->
<!-- LINE NUMBERS ARE GIVEN BELOW. -->

#### 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)..

<a id="sphx-glr-auto-examples-unfolded-demo-unfolded-constrained-lista-py"></a>

# Unfolded Chambolle-Pock for constrained image inpainting

Image inpainting consists in solving $y = Ax$ where $A$ is a mask operator.
This problem can be reformulated as the following minimization problem:

$$
\begin{equation*}
\underset{x}{\operatorname{min}} \,\, \iota_{\mathcal{B}_2(y, r)}(Ax) + \regname(x)
\end{equation*}
$$

where $\iota_{\mathcal{B}_2(y, r)}$ is the indicator function of the ball of radius $r$ centered at
$y$ for the $\ell_2$ norm, and $\regname$ is a regularisation. Recall that the indicator function of
a convex set $\mathcal{C}$ is defined as $\iota_{\mathcal{C}}(x) = 0$ if $x \in \mathcal{C}$ and
$\iota_{\mathcal{C}}(x) = +\infty$ otherwise.

In this example, we unfold the Chambolle-Pock algorithm to solve this problem, and learn the thresholding parameters of
a wavelet denoiser in a LISTA fashion.

<!-- GENERATED FROM PYTHON SOURCE LINES 24-36 -->
```Python
from pathlib import Path
import torch
from torch.utils.data import DataLoader
from torchvision import transforms
import deepinv as dinv
from deepinv.utils import load_dataset
from deepinv.optim.data_fidelity import IndicatorL2
from deepinv.optim.prior import PnP
from deepinv.optim import PDCP
```

<!-- GENERATED FROM PYTHON SOURCE LINES 38-41 -->

## Setup paths for data loading and results.

<!-- GENERATED FROM PYTHON SOURCE LINES 41-51 -->
```Python
BASE_DIR = Path(".")
DATA_DIR = BASE_DIR / "measurements"
CKPT_DIR = BASE_DIR / "ckpts"

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

device = dinv.utils.get_device()
```

```none
Selected GPU 0 with 29513.25 MiB free memory
```

<!-- GENERATED FROM PYTHON SOURCE LINES 52-56 -->

## Load base image datasets and degradation operators.

In this example, we use the CBSD68 dataset for training and the set3c dataset for testing.
We work with images of size 32x32 if no GPU is available, else 128x128.

<!-- GENERATED FROM PYTHON SOURCE LINES 56-74 -->
```Python
operation = "inpainting"
train_dataset_name = "CBSD68"
test_dataset_name = "set3c"
img_size = 128 if torch.cuda.is_available() else 32

test_transform = transforms.Compose(
    [transforms.CenterCrop(img_size), transforms.ToTensor()]
)
train_transform = transforms.Compose(
    [transforms.RandomCrop(img_size), transforms.ToTensor()]
)

train_base_dataset = load_dataset(train_dataset_name, transform=train_transform)
test_base_dataset = load_dataset(test_dataset_name, transform=test_transform)
```

<!-- GENERATED FROM PYTHON SOURCE LINES 75-83 -->

## Define forward operator and generate dataset

We define an inpainting operator that randomly masks pixels with probability 0.5.

A dataset of pairs of measurements and ground truth images is then generated using the
[`deepinv.datasets.generate_dataset()`](https://deepinv.org/api/stubs/deepinv.datasets.generate_dataset.html.md#deepinv.datasets.generate_dataset) function.

Once the dataset is generated, we can load it using the [`deepinv.datasets.HDF5Dataset`](https://deepinv.org/api/stubs/deepinv.datasets.HDF5Dataset.html.md#deepinv.datasets.HDF5Dataset) class.

<!-- GENERATED FROM PYTHON SOURCE LINES 83-130 -->
```Python
n_channels = 3  # 3 for color images, 1 for gray-scale images
probability_mask = 0.5  # probability to mask pixel

# Generate inpainting operator
physics = dinv.physics.Inpainting(
    img_size=(n_channels, img_size, img_size), mask=probability_mask, device=device
)


# Use parallel dataloader if using a GPU to speed up training,
# otherwise, as all computes are on CPU, use synchronous data loading.
num_workers = 4 if torch.cuda.is_available() else 0
n_images_max = (
    100 if torch.cuda.is_available() else 50
)  # maximal number of images used for training
my_dataset_name = "demo_training_inpainting"
measurement_dir = DATA_DIR / train_dataset_name / operation
deepinv_datasets_path = dinv.datasets.generate_dataset(
    train_dataset=train_base_dataset,
    test_dataset=test_base_dataset,
    physics=physics,
    device=device,
    save_dir=measurement_dir,
    train_datapoints=n_images_max,
    num_workers=num_workers,
    dataset_filename=str(my_dataset_name),
)

train_dataset = dinv.datasets.HDF5Dataset(
    path=deepinv_datasets_path, train=True, use_dict_output=True
)
test_dataset = dinv.datasets.HDF5Dataset(
    path=deepinv_datasets_path, train=False, use_dict_output=True
)


train_batch_size = 32 if torch.cuda.is_available() else 3
test_batch_size = 32 if torch.cuda.is_available() else 3

train_dataloader = DataLoader(
    train_dataset, batch_size=train_batch_size, num_workers=num_workers, shuffle=True
)
test_dataloader = DataLoader(
    test_dataset, batch_size=test_batch_size, num_workers=num_workers, shuffle=False
)
```

```none
Dataset has been saved at measurements/CBSD68/inpainting/demo_training_inpainting0.h5
```

<!-- GENERATED FROM PYTHON SOURCE LINES 131-143 -->

## Set up the reconstruction network

We unfold the Chambolle-Pock algorithm as follows:

> $$
> u_{k+1} &= \operatorname{prox}_{\sigma d^*}(u_k + \sigma A z_k) \\
> x_{k+1} &= \operatorname{D_{\sigma}}(x_k-\tau A^\top u_{k+1}) \\
> z_{k+1} &= 2x_{k+1} -x_k \\

> $$

where $\operatorname{D_{\sigma}}$ is a wavelet denoiser with thresholding parameters $\sigma$.

The learnable parameters of our network are $\tau$ and $\sigma$.

<!-- GENERATED FROM PYTHON SOURCE LINES 143-186 -->
```Python
# Select the data fidelity term
data_fidelity = IndicatorL2(radius=0.0)

# Set up the trainable denoising prior; here, the soft-threshold in a wavelet basis.
# If the prior is initialized with a list of length max_iter,
# then a distinct weight is trained for each CP iteration.
# For fixed trained model prior across iterations, initialize with a single model.
max_iter = 30 if torch.cuda.is_available() else 20  # Number of unrolled iterations
level = 3
prior = [
    PnP(denoiser=dinv.models.WaveletDenoiser(wv="db8", level=level, device=device))
    for i in range(max_iter)
]

# Unrolled optimization algorithm parameters
stepsize = [
    1.0
] * max_iter  # initialization of the stepsizes. A distinct stepsize is trained for each iteration.
sigma_denoiser = [
    0.01 * torch.ones(1, level, 3)
] * max_iter  # thresholding parameter of the wavelet denoiser.

stepsize_dual = 1.0  # dual stepsize for Chambolle-Pock


# define which parameters are trainable : here, both the regularization parameters and the primal/dual stepsizes are learned.
trainable_params = ["sigma_denoiser", "stepsize", "stepsize_dual"]

# Define the unfolded trainable model. # See the documentation of the Primal Dual CP algorithm :class:`deepinv.optim.PDCP` for more details.
model = PDCP(
    stepsize=stepsize,
    sigma_denoiser=sigma_denoiser,
    stepsize_dual=stepsize_dual,
    K=physics.A,
    K_adjoint=physics.A_adjoint,
    trainable_params=trainable_params,
    data_fidelity=data_fidelity,
    max_iter=max_iter,
    prior=prior,
    unfold=True,
)
```

<!-- GENERATED FROM PYTHON SOURCE LINES 187-199 -->

## Train the model

We train the model using the [`deepinv.Trainer`](https://deepinv.org/api/stubs/deepinv.Trainer.html.md#deepinv.Trainer) class.

We perform supervised learning and use the mean squared error as loss function. This can be easily done using the
[`deepinv.loss.SupLoss`](https://deepinv.org/api/stubs/deepinv.loss.SupLoss.html.md#deepinv.loss.SupLoss) class.

#### NOTE
In this example, we only train for a few epochs to keep the training time short on CPU.
For a good reconstruction quality, we recommend to train for at least 50 epochs.

<!-- GENERATED FROM PYTHON SOURCE LINES 199-228 -->
```Python
epochs = 10 if torch.cuda.is_available() else 5  # choose training epochs
learning_rate = 1e-3

verbose = True  # print training information

# choose training losses
losses = dinv.loss.SupLoss(metric=dinv.metric.MSE())

# choose optimizer and scheduler
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=1e-8)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=int(epochs * 0.8))

trainer = dinv.Trainer(
    model=model,
    scheduler=scheduler,
    losses=losses,
    device=device,
    optimizer=optimizer,
    physics=physics,
    train_dataloader=train_dataloader,
    save_path=str(CKPT_DIR / operation),
    verbose=verbose,
    show_progress_bar=False,  # disable progress bar for better vis in sphinx gallery.
    epochs=epochs,
)

model = trainer.train()
```

```none
/local/jtachell/deepinv/deepinv/deepinv/training/trainer.py:1415: UserWarning: non_blocking_transfers=True but DataLoader.pin_memory=False; set pin_memory=True to overlap host-device copies with compute.
  self.setup_train()
The model has 301 trainable parameters
Train epoch 0: TotalLoss=0.019, PSNR=20.959
Train epoch 1: TotalLoss=0.003, PSNR=26.244
Train epoch 2: TotalLoss=0.003, PSNR=26.485
Train epoch 3: TotalLoss=0.003, PSNR=26.627
Train epoch 4: TotalLoss=0.003, PSNR=26.74
Train epoch 5: TotalLoss=0.003, PSNR=26.833
Train epoch 6: TotalLoss=0.003, PSNR=26.912
Train epoch 7: TotalLoss=0.003, PSNR=26.98
Train epoch 8: TotalLoss=0.003, PSNR=27.028
Train epoch 9: TotalLoss=0.003, PSNR=27.034
```

<!-- GENERATED FROM PYTHON SOURCE LINES 229-234 -->

## Test the network

We can now test the trained network using the [`deepinv.test()`](https://deepinv.org/api/stubs/deepinv.test.html.md#deepinv.test) function.

The testing function will compute test_psnr metrics and plot and save the results.

<!-- GENERATED FROM PYTHON SOURCE LINES 234-240 -->
```Python
plot_images = True
method = "artifact_removal"

trainer.test(test_dataloader)
```

```none
/local/jtachell/deepinv/deepinv/deepinv/training/trainer.py:1611: UserWarning: non_blocking_transfers=True but DataLoader.pin_memory=False; set pin_memory=True to overlap host-device copies with compute.
  self.setup_train(train=False)
Eval epoch 0: PSNR=24.736, PSNR no learning=6.654
Test results:
PSNR no learning: 6.654 +- 0.894
PSNR: 24.736 +- 0.326

{'PSNR no learning': 6.6535491943359375, 'PSNR no learning_std': 0.8941914593348746, 'PSNR': 24.73596445719401, 'PSNR_std': 0.326294800623154, 'runtime': 0.060470234137028456, 'peak_gpu_memory_usage': 0.7331256866455078}
```

<!-- GENERATED FROM PYTHON SOURCE LINES 241-244 -->

## Saving the model

We can save the trained model following the standard PyTorch procedure.

<!-- GENERATED FROM PYTHON SOURCE LINES 244-248 -->
```Python
# Save the model
torch.save(model.state_dict(), CKPT_DIR / operation / "model.pth")
```

<!-- GENERATED FROM PYTHON SOURCE LINES 249-253 -->

## Loading the model

Similarly, we can load our trained unfolded architecture following the standard PyTorch procedure.
To check that the loading is performed correctly, we use new variables for the initialization of the model.

<!-- GENERATED FROM PYTHON SOURCE LINES 253-317 -->
```Python
# Set up the trainable denoising prior; here, the soft-threshold in a wavelet basis.
level = 3
model_spec = {
    "name": "waveletprior",
    "args": {"wv": "db8", "level": level, "device": device},
}
# If the prior is initialized with a list of length max_iter,
# then a distinct weight is trained for each PGD iteration.
# For fixed trained model prior across iterations, initialize with a single model.
max_iter = 30 if torch.cuda.is_available() else 20  # Number of unrolled iterations
prior_new = [
    PnP(denoiser=dinv.models.WaveletDenoiser(wv="db8", level=level, device=device))
    for i in range(max_iter)
]

# Unrolled optimization algorithm parameters
stepsize = [
    1.0
] * max_iter  # initialization of the stepsizes. A distinct stepsize is trained for each iteration.
sigma_denoiser = [0.01 * torch.ones(1, level, 3)] * max_iter
stepsize_dual = 1.0  # stepsize for Chambolle-Pock


model_new = PDCP(
    unfold=True,
    stepsize=stepsize,
    sigma_denoiser=sigma_denoiser,
    stepsize_dual=stepsize_dual,
    K=physics.A,
    K_adjoint=physics.A_adjoint,
    trainable_params=trainable_params,
    data_fidelity=data_fidelity,
    max_iter=max_iter,
    prior=prior_new,
    g_first=False,
)

model_new.load_state_dict(torch.load(CKPT_DIR / operation / "model.pth"))
model_new.eval()


# Test the model and check that the results are the same as before saving.
dinv.training.test(
    model_new, test_dataloader, physics=physics, device=device, show_progress_bar=False
)

# Plot the results
test_sample = next(iter(test_dataloader))["x"]
model.eval()
test_sample = test_sample.to(device)

# Get the measurements and the ground truth
y = physics(test_sample)
with torch.no_grad():
    rec = model(y, physics=physics)

backprojected = physics.A_adjoint(y)

dinv.utils.plot(
    [backprojected, rec, test_sample],
    titles=["Linear", "Reconstruction", "Ground truth"],
    suptitle="Reconstruction results",
)
```

```none
Eval epoch 0: PSNR=24.736, PSNR no learning=6.654
Test results:
PSNR no learning: 6.654 +- 0.894
PSNR: 24.736 +- 0.326
/local/jtachell/deepinv/deepinv/deepinv/utils/plotting.py:519: UserWarning: This figure was using a layout engine that is incompatible with subplots_adjust and/or tight_layout; not calling subplots_adjust.
  fig.subplots_adjust(top=0.75)
```

**Total running time of the script:** (0 minutes 12.983 seconds)

🏷 Tags: Inpainting

<a id="sphx-glr-download-auto-examples-unfolded-demo-unfolded-constrained-lista-py"></a>

[`Download Jupyter notebook: demo_unfolded_constrained_LISTA.ipynb`](https://deepinv.org/auto_examples/unfolded/_downloads/8951ca8fee4e9b4ca466202eee01392f/demo_unfolded_constrained_LISTA.ipynb)

[`Download Python source code: demo_unfolded_constrained_LISTA.py`](https://deepinv.org/auto_examples/unfolded/_downloads/54991abcb2105cbefad71602185c6e3b/demo_unfolded_constrained_LISTA.py)

[`Download zipped: demo_unfolded_constrained_LISTA.zip`](https://deepinv.org/auto_examples/unfolded/_downloads/0f9ad12bd11e94f8fc4c5d9a52b9ed20/demo_unfolded_constrained_LISTA.zip)

[Gallery generated by Sphinx-Gallery](https://sphinx-gallery.github.io)