#### 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 learning with measurement splitting
We demonstrate self-supervised learning with measurement splitting, to
train a denoiser network on the MNIST dataset. The physics here is noisy
computed tomography, as is the case in Noise2Inverse Hendriksen *et al.*[1](#footcite-hendriksen2020noise2inverse). Note this example
can also be easily applied to undersampled multicoil MRI as is the case
in SSDU Yaman *et al.*[2](#footcite-yaman2020self).
Measurement splitting constructs a ground-truth free loss
$\frac{m}{m_2}\| y_2 - A_2 \inversef{y_1}{A_1}\|^2$ by splitting
the measurement and the forward operator using a randomly generated
mask.
See [`deepinv.loss.SplittingLoss`](https://deepinv.org/api/stubs/deepinv.loss.SplittingLoss.html.md#deepinv.loss.SplittingLoss) for full details.
```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
torch.manual_seed(0)
device = dinv.utils.get_device()
BASE_DIR = Path(".")
DATA_DIR = BASE_DIR / "measurements"
ORIGINAL_DATA_HOME = get_cache_home() / "datasets" / "MNIST"
```
```none
Selected GPU 0 with 17549.25 MiB free memory
```
## Define loss
Our implementation has multiple optional parameters that control how the
splitting is to be achieved. For example, you can:
- Use `split_ratio` to set the ratio of pixels used in the forward
pass vs the loss;
- Define custom masking methods using a `mask_generator` such as
[`deepinv.physics.generator.BernoulliSplittingMaskGenerator`](https://deepinv.org/api/stubs/deepinv.physics.generator.BernoulliSplittingMaskGenerator.html.md#deepinv.physics.generator.BernoulliSplittingMaskGenerator)
or [`deepinv.physics.generator.GaussianSplittingMaskGenerator`](https://deepinv.org/api/stubs/deepinv.physics.generator.GaussianSplittingMaskGenerator.html.md#deepinv.physics.generator.GaussianSplittingMaskGenerator);
- Use `eval_n_samples` to set how many realisations of the random
mask is used at evaluation time;
- Optionally disable measurement splitting at evaluation time using
`eval_split_input` (as is the case in SSDU Yaman *et al.*[2](#footcite-yaman2020self)).
- Average over both input and output masks at evaluation time using
`eval_split_output`. See [`deepinv.loss.SplittingLoss`](https://deepinv.org/api/stubs/deepinv.loss.SplittingLoss.html.md#deepinv.loss.SplittingLoss) for
details.
Note that after the model has been defined, the loss must also “adapt”
the model.
```Python
loss = dinv.loss.SplittingLoss(split_ratio=0.6, eval_split_input=True, eval_n_samples=5)
```
## Prepare data
We use the `torchvision` MNIST dataset, and use noisy tomography
physics (with number of angles equal to the image size) for the forward
operator.
#### 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 `train_datapoints=test_datapoints=None` to get the best results.
```Python
transform = transforms.Compose([transforms.ToTensor()])
train_dataset = datasets.MNIST(
root=ORIGINAL_DATA_HOME, train=True, transform=transform, download=True
)
test_dataset = datasets.MNIST(
root=ORIGINAL_DATA_HOME, train=False, transform=transform, download=True
)
physics = dinv.physics.Tomography(
angles=28,
img_width=28,
noise_model=dinv.physics.noise.GaussianNoise(0.1),
device=device,
)
deepinv_datasets_path = dinv.datasets.generate_dataset(
train_dataset=train_dataset,
test_dataset=test_dataset,
physics=physics,
device=device,
save_dir=DATA_DIR,
train_datapoints=100,
test_datapoints=10,
)
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_dataloader = DataLoader(train_dataset, shuffle=True)
test_dataloader = DataLoader(test_dataset, shuffle=False)
```
```none
/local/jtachell/deepinv/deepinv/deepinv/physics/tomography.py:186: UserWarning: The default value of `normalize` is not specified and will be automatically set to `True`. Set `normalize` explicitly to `True` or `False` to avoid this warning.
warn(
Dataset has been saved at measurements/dinv_dataset0.h5
```
## Define model
We use a simple U-Net architecture with 2 scales as the denoiser
network.
To reduce training time, we use a pretrained model. Here we demonstrate
training with 100 images for 1 epoch, after having loaded a pretrained
model trained that was with 1000 images for 20 epochs.
#### NOTE
When using the splitting loss, the model must be “adapted” by the loss, as its forward pass takes only a subset of the pixels, not the full image.
```Python
model = dinv.models.ArtifactRemoval(
dinv.models.UNet(in_channels=1, out_channels=1, scales=2).to(device), mode="pinv"
)
model = loss.adapt_model(model)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-8)
# Load pretrained model
file_name = "demo_measplit_mnist_tomography.pth"
url = get_weights_url(model_name="measplit", 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"])
```
### Train and test 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
trainer = dinv.Trainer(
model=model,
physics=physics,
epochs=1,
losses=loss,
optimizer=optimizer,
device=device,
train_dataloader=train_dataloader,
eval_dataloader=test_dataloader,
metrics=None, # no supervised metrics
early_stop=2, # we can use early stopping as we have a validation loss
compute_eval_losses=True, # use self-supervised loss for evaluation
early_stop_on_losses=True, # stop using self-supervised eval loss
plot_images=False,
save_path=None,
verbose=True,
show_progress_bar=False,
no_learning_method="A_dagger", # use pseudo-inverse as no-learning baseline
)
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
/local/jtachell/deepinv/deepinv/deepinv/loss/measplit.py:275: UserWarning: Mask generator not defined. Using new Bernoulli mask generator with shape torch.Size([40, 28]).
warn(
/local/jtachell/deepinv/deepinv/deepinv/physics/forward.py:1020: UserWarning: At least one input physics is a DecomposablePhysics, but resulting physics will not be decomposable. `A_dagger` and `prox_l2` will fall back to approximate methods, which may impact performance.
warnings.warn(
Train epoch 0: TotalLoss=0.034
Eval epoch 0: TotalLoss=0.015
Best model saved at epoch 1
```
Test and visualize the model outputs using a small test set. We set the
output to average over 5 iterations of random mask realisations. The
trained model improves on the no-learning reconstruction by ~7dB.
```Python
trainer.plot_images = True
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.015, PSNR=6.443, PSNR no learning=-7.519
Test results:
PSNR no learning: -7.519 +- 0.393
PSNR: 6.443 +- 0.618
{'PSNR no learning': -7.5194176197052, 'PSNR no learning_std': 0.3926154011446893, 'PSNR': 6.443156433105469, 'PSNR_std': 0.6182282766836074, 'runtime': 2.14781123585999, 'peak_gpu_memory_usage': 0.7409062385559082}
```
Demonstrate the effect of not averaging over multiple realisations of
the splitting mask at evaluation time, by setting `eval_n_samples=1`.
We have a worse performance:
```Python
model.eval_n_samples = 1
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)
/local/jtachell/deepinv/deepinv/deepinv/physics/forward.py:1020: UserWarning: At least one input physics is a DecomposablePhysics, but resulting physics will not be decomposable. `A_dagger` and `prox_l2` will fall back to approximate methods, which may impact performance.
warnings.warn(
Eval epoch 0: TotalLoss=0.015, PSNR=2.954, PSNR no learning=-7.537
Test results:
PSNR no learning: -7.537 +- 0.396
PSNR: 2.954 +- 0.691
{'PSNR no learning': -7.536773586273194, 'PSNR no learning_std': 0.39602266136826014, 'PSNR': 2.953848588466644, 'PSNR_std': 0.6911939444599268, 'runtime': 0.9846608899533749, 'peak_gpu_memory_usage': 0.7408857345581055}
```
Furthermore, we can disable measurement splitting at evaluation
altogether by setting `eval_split_input` to False (this is done in
SSDU Yaman *et al.*[2](#footcite-yaman2020self)). This generally is
worse than MC averaging:
```Python
model.eval_split_input = False
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)
/local/jtachell/deepinv/deepinv/deepinv/physics/forward.py:1020: UserWarning: At least one input physics is a DecomposablePhysics, but resulting physics will not be decomposable. `A_dagger` and `prox_l2` will fall back to approximate methods, which may impact performance.
warnings.warn(
Eval epoch 0: TotalLoss=0.015, PSNR=0.301, PSNR no learning=-7.516
Test results:
PSNR no learning: -7.516 +- 0.397
PSNR: 0.301 +- 0.472
{'PSNR no learning': -7.5157218933105465, 'PSNR no learning_std': 0.3970079276348475, 'PSNR': 0.30101901553571225, 'PSNR_std': 0.4723353251216, 'runtime': 0.9194617699831724, 'peak_gpu_memory_usage': 0.7408766746520996}
```
* **References:**
* Allard Adriaan Hendriksen, Daniël Maria Pelt, and K Joost Batenburg. Noise2inverse: self-supervised deep convolutional denoising for tomography. *IEEE Transactions on Computational Imaging*, 6:1320–1335, 2020.
* Burhaneddin Yaman, Seyed Amir Hossein Hosseini, Steen Moeller, Jutta Ellermann, Kâmil Uğurbil, and Mehmet Akçakaya. Self-supervised learning of physics-guided reconstruction neural networks without fully sampled reference data. *Magnetic resonance in medicine*, 84(6):3172–3191, 2020.
**Total running time of the script:** (0 minutes 8.187 seconds)
🏷 Tags: Denoising
[`Download Jupyter notebook: demo_splitting_loss.ipynb`](https://deepinv.org/auto_examples/self-supervised-learning/_downloads/9df1cf2e94dc0725a73bf7d3d2a8ee90/demo_splitting_loss.ipynb)
[`Download Python source code: demo_splitting_loss.py`](https://deepinv.org/auto_examples/self-supervised-learning/_downloads/32412a422ce6ce1fd4c9f16c4956ac91/demo_splitting_loss.py)
[`Download zipped: demo_splitting_loss.zip`](https://deepinv.org/auto_examples/self-supervised-learning/_downloads/9a0b17eec4328531a061fa837f14626e/demo_splitting_loss.zip)
[Gallery generated by Sphinx-Gallery](https://sphinx-gallery.github.io)