<!-- DO NOT EDIT. -->
<!-- THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. -->
<!-- TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: -->
<!-- "auto_examples/unfolded/demo_vanilla_unfolded.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-vanilla-unfolded-py"></a>

# Vanilla Unfolded algorithm for super-resolution

This is a simple example to show how to use vanilla unfolded Plug-and-Play.
The DnCNN denoiser and the algorithm parameters (stepsize, regularization parameters) are trained jointly.
For simplicity, we show how to train the algorithm on a  small dataset. For optimal results, use a larger dataset.

<!-- GENERATED FROM PYTHON SOURCE LINES 11-24 -->
```Python
import deepinv as dinv
import torch
from deepinv.models.utils import get_weights_url
from torch.utils.data import DataLoader
from deepinv.optim.data_fidelity import L2
from deepinv.optim.prior import PnP
from deepinv.optim import DRS
from torchvision import transforms
from deepinv.utils import get_cache_home
from deepinv.datasets import BSDS500
```

<!-- GENERATED FROM PYTHON SOURCE LINES 26-29 -->

## Setup paths for data loading and results.

<!-- GENERATED FROM PYTHON SOURCE LINES 29-40 -->
```Python
BASE_DIR = get_cache_home() / "demo_unfolded_sr"
DATA_DIR = BASE_DIR / "measurements"
RESULTS_DIR = BASE_DIR / "results"
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 27559.25 MiB free memory
```

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

## Load base image datasets and degradation operators.

In this example, we use the CBSD500 dataset for training and the Set3C dataset for testing.

<!-- GENERATED FROM PYTHON SOURCE LINES 44-49 -->
```Python
img_size = 64 if torch.cuda.is_available() else 32
n_channels = 3  # 3 for color images, 1 for gray-scale images
operation = "super-resolution"
```

<!-- GENERATED FROM PYTHON SOURCE LINES 50-53 -->

## Generate a dataset of low resolution images and load it.

We use the Downsampling class from the physics module to generate a dataset of low resolution images.

<!-- GENERATED FROM PYTHON SOURCE LINES 53-107 -->
```Python
# For simplicity, we use a small dataset for training.
# To be replaced for optimal results. For example, you can use the larger DIV2K or LSDIR datasets (also provided in the library).

# Specify the  train and test transforms to be applied to the input images.
test_transform = transforms.Compose(
    [transforms.CenterCrop(img_size), transforms.ToTensor()]
)
train_transform = transforms.Compose(
    [transforms.RandomCrop(img_size), transforms.ToTensor()]
)
# Define the base train and test datasets of clean images.
train_base_dataset = BSDS500(download=True, train=True, transform=train_transform)
test_base_dataset = BSDS500(download=False, train=False, transform=test_transform)

# Use parallel dataloader if using a GPU to speed up training, otherwise, as all computes are on CPU, use synchronous
# dataloading.
num_workers = 4 if torch.cuda.is_available() else 0

# Degradation parameters
factor = 2
noise_level_img = 0.0075

# Generate the gaussian blur downsampling operator.
physics = dinv.physics.Downsampling(
    filter="gaussian",
    img_size=(n_channels, img_size, img_size),
    factor=factor,
    device=device,
    noise_model=dinv.physics.GaussianNoise(sigma=noise_level_img),
)
my_dataset_name = "demo_unfolded_sr"
n_images_max = (
    None if torch.cuda.is_available() else 10
)  # max number of images used for training (use all if you have a GPU)
measurement_dir = DATA_DIR / "BSDS500" / operation
generated_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=generated_datasets_path, train=True, use_dict_output=True
)
test_dataset = dinv.datasets.HDF5Dataset(
    path=generated_datasets_path, train=False, use_dict_output=True
)
```

```none
/local/jtachell/deepinv/deepinv/deepinv/datasets/bsds500.py:56: DeprecationWarning: The tuple format for dataset outputs is deprecated and will be removed in a future version.It is recommended to set `use_dict_output=True` for better readability and flexibility in returned outputs.The default is currently `False` for backward compatibility, but will be switched to `True` in a future version.
  super().__init__(use_dict_output=use_dict_output)
/local/jtachell/deepinv/deepinv/deepinv/datasets/datagenerator.py:624: UserWarning: Dataset /local/jtachell/.cache/deepinv/demo_unfolded_sr/measurements/BSDS500/super-resolution/demo_unfolded_sr0.h5 already exists, this will close and overwrite the previous dataset.
  warn(
/local/jtachell/deepinv/deepinv/.pixi/envs/docs/lib/python3.12/multiprocessing/popen_fork.py:66: DeprecationWarning: This process (pid=4088353) is multi-threaded, use of fork() may lead to deadlocks in the child.
  self.pid = os.fork()
Dataset has been saved at /local/jtachell/.cache/deepinv/demo_unfolded_sr/measurements/BSDS500/super-resolution/demo_unfolded_sr0.h5
```

<!-- GENERATED FROM PYTHON SOURCE LINES 108-118 -->

## Define the unfolded PnP algorithm.

The chosen algorithm is here DRS (Douglas-Rachford Splitting).
Note that if the prior (resp. a parameter) is initialized with a list of length max_iter,
then a distinct model (resp. parameter) is trained for each iteration.
For fixed trained model prior (resp. parameter) across iterations, initialize with a single element.

#### TIP
While in theory any initial stepsize could be adapted during training, in practice the best performance
is obtained by starting with a stepsize close to the largest stable one, i.e. $1/\|A\|^2$

<!-- GENERATED FROM PYTHON SOURCE LINES 118-159 -->
```Python
# Unrolled optimization algorithm parameters
max_iter = 5  # number of unfolded layers

# Select the data fidelity term
data_fidelity = L2()

# Set up the trainable denoising prior
# Here the prior model is common for all iterations
prior = PnP(denoiser=dinv.models.DnCNN(depth=20, pretrained="download").to(device))

# The parameters are initialized with a list of length max_iter, so that a distinct parameter is trained for each iteration.
norm = physics.compute_norm(
    torch.ones((1, 3, img_size, img_size), device=device), squared=True
)
stepsize = [1.0 / norm] * max_iter  # stepsize of the algorithm
sigma_denoiser = [
    1.0
] * max_iter  # noise level parameter of the denoiser (not used by DnCNN)
beta = 1.0  # relaxation parameter of the Douglas-Rachford splitting
trainable_params = [
    "stepsize",
    "beta",
    "sigma_denoiser",
]  # define which parameters are trainable

# Logging parameters
verbose = True

# Define the unfolded trainable model.
model = DRS(
    stepsize=stepsize,
    sigma_denoiser=sigma_denoiser,
    beta=beta,
    trainable_params=trainable_params,
    data_fidelity=data_fidelity,
    max_iter=max_iter,
    prior=prior,
    unfold=True,
)
```

```none
/local/jtachell/deepinv/deepinv/deepinv/physics/forward.py:639: DeprecationWarning: Using `compute_norm(squared=True)` is deprecated. Use `compute_sqnorm()` instead to compute the squared spectral norm (||A^T A||_2). In a future version, `compute_norm()` will compute the non-squared spectral norm (||A||_2) by default.
  warnings.warn(
Power iteration converged at iteration 13, ||A^T A||_2=0.24
```

<!-- GENERATED FROM PYTHON SOURCE LINES 160-167 -->

## Define the training parameters.

We use the Adam optimizer and the StepLR scheduler.

#### NOTE
We use a pretrained checkpoint here (trained for 30 epochs) to speed up the demo,
but you can reproduce results by training for the whole 30 epochs from scratch.

<!-- GENERATED FROM PYTHON SOURCE LINES 167-197 -->
```Python
# training parameters
epochs = 3 if torch.cuda.is_available() else 1
learning_rate = 5e-4
train_batch_size = 32 if torch.cuda.is_available() else 1
test_batch_size = 3

# choose optimizer and scheduler
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=1e-8)

# start with a pretrained model to reduce training time
# comment out if you wan't to train from scratch
file_name = "demo_vanilla_unfolded.pth"
url = get_weights_url(model_name="demo", file_name=file_name)
ckpt = torch.hub.load_state_dict_from_url(
    url, map_location=lambda storage, loc: storage, file_name=file_name
)
model.load_state_dict(ckpt["state_dict"])
optimizer.load_state_dict(ckpt["optimizer"])

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

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
)
```

<!-- GENERATED FROM PYTHON SOURCE LINES 198-201 -->

## Train the network

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

<!-- GENERATED FROM PYTHON SOURCE LINES 201-220 -->
```Python
trainer = dinv.Trainer(
    model,
    physics=physics,
    train_dataloader=train_dataloader,
    eval_dataloader=test_dataloader,
    epochs=epochs,
    losses=losses,
    optimizer=optimizer,
    device=device,
    early_stop=True,  # set to None to disable early stopping
    save_path=str(CKPT_DIR / operation),
    verbose=verbose,
    show_progress_bar=False,  # disable progress bar for better vis in sphinx gallery.
)

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 668238 trainable parameters
/local/jtachell/deepinv/deepinv/deepinv/training/trainer.py:565: UserWarning: early_stop should be an integer or None. Setting early_stop=3. This behaviour will be deprecated in future versions.
  warnings.warn(
Train epoch 0: TotalLoss=0.003, PSNR=27.91
Eval epoch 0: PSNR=25.152
Best model saved at epoch 1
Train epoch 1: TotalLoss=0.003, PSNR=27.955
Eval epoch 1: PSNR=25.219
Best model saved at epoch 2
Train epoch 2: TotalLoss=0.003, PSNR=27.996
Eval epoch 2: PSNR=25.143
```

<!-- GENERATED FROM PYTHON SOURCE LINES 221-225 -->

## Test the network

<!-- GENERATED FROM PYTHON SOURCE LINES 225-243 -->
```Python
trainer.test(test_dataloader)

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
/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=25.143, PSNR no learning=9.123
Test results:
PSNR no learning: 9.123 +- 2.904
PSNR: 25.143 +- 4.099
/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 31.897 seconds)

🏷 Tags: Super-resolution

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

[`Download Jupyter notebook: demo_vanilla_unfolded.ipynb`](https://deepinv.org/auto_examples/unfolded/_downloads/d413f127f8fcdc4f25ebe8912dc1f800/demo_vanilla_unfolded.ipynb)

[`Download Python source code: demo_vanilla_unfolded.py`](https://deepinv.org/auto_examples/unfolded/_downloads/aa635a7f9321b9a4a94be443b1752796/demo_vanilla_unfolded.py)

[`Download zipped: demo_vanilla_unfolded.zip`](https://deepinv.org/auto_examples/unfolded/_downloads/5af3f2cac86b04c291e6c72d72b7d37a/demo_vanilla_unfolded.zip)

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