#### 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)..
# Self-supervised denoising with the Generalized R2R loss.
This example shows you how to train a denoiser network in a fully self-supervised way,
using noisy images only via the Generalized Recorrupted2Recorrupted (GR2R) loss Monroy *et al.*[1](#footcite-monroy2025generalized),
which exploits knowledge about the noise distribution. You can change the noise distribution by selecting
from predefined noise models such as Gaussian, Poisson, and Gamma noise.
```Python
from pathlib import Path
import torch
from torch.utils.data import DataLoader
from torchvision import transforms, datasets
import deepinv as dinv
from deepinv.utils import get_cache_home
from deepinv.models.utils import get_weights_url
```
## Setup paths for data loading and results.
```Python
BASE_DIR = Path(".")
DATA_DIR = BASE_DIR / "measurements"
CKPT_DIR = BASE_DIR / "ckpts"
ORIGINAL_DATA_DIR = get_cache_home() / "datasets" / "MNIST"
# Set the global random seed from pytorch to ensure reproducibility of the example.
torch.manual_seed(0)
device = dinv.utils.get_device()
print(device)
```
```none
Selected GPU 0 with 22453.25 MiB free memory
cuda:0
```
## Load base image datasets
In this example, we use the MNIST dataset as the base image dataset.
```Python
operation = "denoising"
train_dataset_name = "MNIST"
transform = transforms.Compose([transforms.ToTensor()])
train_dataset = datasets.MNIST(
root=ORIGINAL_DATA_DIR, train=True, transform=transform, download=True
)
test_dataset = datasets.MNIST(
root=ORIGINAL_DATA_DIR, train=False, transform=transform, download=True
)
```
## Generate a dataset of noisy images
Generate a dataset of noisy images corrupted by Poisson noise.
The predefined noise models in the physics module include Gaussian, Poisson, and Gamma noise.
Here, we use Poisson noise as an example, but you can also use Gaussian or Gamma noise.
#### NOTE
We use a subset of the whole training set to reduce the computational load of the example.
We recommend to use the whole set by setting `n_images_max=None` to get the best results.
```Python
# defined physics
predefined_noise_models = dict(
gaussian=dinv.physics.GaussianNoise(sigma=0.1),
poisson=dinv.physics.PoissonNoise(gain=0.5),
gamma=dinv.physics.GammaNoise(l=10.0),
)
noise_name = "poisson"
noise_model = predefined_noise_models[noise_name]
physics = dinv.physics.Denoising(noise_model)
operation = f"{operation}_{noise_name}"
# Use parallel dataloader if using a GPU to speed up training,
# otherwise, as all computes are on CPU, use synchronous data loading.
num_workers = 0 if torch.cuda.is_available() else 0
n_images_max = (
100 if torch.cuda.is_available() else 5
) # number of images used for training
measurement_dir = DATA_DIR / train_dataset_name / operation
deepinv_datasets_path = dinv.datasets.generate_dataset(
train_dataset=train_dataset,
test_dataset=test_dataset,
physics=physics,
device=device,
save_dir=measurement_dir,
train_datapoints=n_images_max,
test_datapoints=n_images_max,
num_workers=num_workers,
dataset_filename="demo_r2r",
)
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
)
```
```none
Dataset has been saved at measurements/MNIST/denoising_poisson/demo_r2r0.h5
```
## Set up the denoiser network
We use a simple U-Net architecture with 2 scales as the denoiser network.
```Python
model = dinv.models.ArtifactRemoval(
dinv.models.UNet(in_channels=1, out_channels=1, scales=2, residual=False).to(device)
)
```
## Set up the training parameters
We set [`deepinv.loss.R2RLoss`](https://deepinv.org/api/stubs/deepinv.loss.R2RLoss.html.md#deepinv.loss.R2RLoss) as the training loss.
#### NOTE
There are GR2R losses for various noise distributions, which can be specified by the noise model.
```Python
epochs = 1 # choose training epochs
learning_rate = 1e-4
batch_size = 64 if torch.cuda.is_available() else 1
# choose self-supervised training loss
loss = dinv.loss.R2RLoss(noise_model=None)
model = loss.adapt_model(model) # important step!
# choose optimizer and scheduler
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate, weight_decay=1e-4)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=int(epochs * 0.8) + 1)
# # start with a pretrained model to reduce training time
if noise_name == "poisson":
file_name = "ckp_10_demo_r2r_poisson.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"])
```
## Train the network
To simulate a realistic self-supervised learning scenario, we do not use any supervised metrics for training,
such as PSNR or SSIM, which require clean ground truth images.
#### TIP
We can use the same self-supervised loss for evaluation, as it does not require clean images,
to monitor the training process (e.g. for early stopping). This is done automatically when `metrics=None` and `early_stop>0` in the trainer.
```Python
verbose = True # print training information
train_dataloader = DataLoader(
train_dataset, batch_size=batch_size, num_workers=num_workers, shuffle=False
)
test_dataloader = DataLoader(
test_dataset, batch_size=batch_size, num_workers=num_workers, shuffle=False
)
# Initialize the trainer
trainer = dinv.Trainer(
model=model,
physics=physics,
epochs=epochs,
scheduler=scheduler,
losses=loss,
optimizer=optimizer,
device=device,
metrics=None, # no supervised metrics
train_dataloader=train_dataloader,
eval_dataloader=test_dataloader,
early_stop=2, # early stop using the self-supervised loss on the test set
compute_eval_losses=True, # use self-supervised loss for evaluation
early_stop_on_losses=True, # stop using self-supervised eval loss
plot_images=True,
save_path=str(CKPT_DIR / operation),
verbose=verbose,
show_progress_bar=False, # disable progress bar for better vis in sphinx gallery.
)
# Train the network
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 444737 trainable parameters
Train epoch 0: TotalLoss=0.418
Eval epoch 0: TotalLoss=0.412
Best model saved at epoch 1
```
## Test the network
We now assume that we have access to a small test set of clean images to evaluate the performance of the trained network.
and we compute the PSNR between the denoised images and the clean ground truth images.
```Python
trainer.test(test_dataloader, metrics=dinv.metric.PSNR())
```
```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: TotalLoss=0.393, PSNR=20.526, PSNR no learning=12.573
Test results:
PSNR no learning: 12.573 +- 1.723
PSNR: 20.526 +- 1.702
{'PSNR no learning': 12.572514038085938, 'PSNR no learning_std': 1.7226337202107147, 'PSNR': 20.526494140625, 'PSNR_std': 1.7016625641964087, 'runtime': 0.061833787243813276, 'peak_gpu_memory_usage': 1.067305564880371}
```
* **References:**
* Brayan Monroy, Jorge Bacca, and Julián Tachella. Generalized recorrupted-to-recorrupted: self-supervised learning beyond gaussian noise. In *Proceedings of the Computer Vision and Pattern Recognition Conference*, 28155–28164. 2025.
**Total running time of the script:** (0 minutes 0.646 seconds)
🏷 Tags: Denoising
[`Download Jupyter notebook: demo_r2r_denoising.ipynb`](https://deepinv.org/auto_examples/self-supervised-learning/_downloads/0dbb9da1987a70e21aae950d9f50f937/demo_r2r_denoising.ipynb)
[`Download Python source code: demo_r2r_denoising.py`](https://deepinv.org/auto_examples/self-supervised-learning/_downloads/576e0efe951302d317a0e6bbc52d3f17/demo_r2r_denoising.py)
[`Download zipped: demo_r2r_denoising.zip`](https://deepinv.org/auto_examples/self-supervised-learning/_downloads/77ee6b7b14f54dd12b67ac5de815c4b7/demo_r2r_denoising.zip)
[Gallery generated by Sphinx-Gallery](https://sphinx-gallery.github.io)