.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "auto_examples/09-learning-workflows/04-annealed-plug-and-play.py" .. LINE NUMBERS ARE GIVEN BELOW. .. only:: html .. note:: :class: sphx-glr-download-link-note :ref:`Go to the end ` to download the full example code. .. rst-class:: sphx-glr-example-title .. _sphx_glr_auto_examples_09-learning-workflows_04-annealed-plug-and-play.py: ====================== Annealed plug-and-play ====================== In lesson 10 you plugged a denoiser someone else had trained, on photographs, into ADMM at one noise level you chose. In this tour you train the denoiser yourself, on MR images alone and without any encoding, and use it as the regularizer of an ADMM reconstruction whose denoising strength decreases over the iterations. You then apply the same network, unchanged, at an acceleration where CG-SENSE breaks down. A plug-and-play reconstruction [#pnp]_ replaces the proximal step of an iterative algorithm by a denoiser. Because the denoiser is trained on images, not on k-space, one network serves every protocol whose images resemble its training images: a change of acceleration, sampling pattern or coil array needs no retraining. The proximal step of a regularizer :math:`\lambda\phi` with ADMM penalty :math:`\rho` is a Gaussian denoiser of noise variance :math:`\sigma^2 = \lambda/\rho`. The first iterates carry the strong incoherent aliasing of the undersampling and the last are nearly consistent with the data, so a noise level that decreases from one to the other [#dpir]_ removes the aliasing first and preserves fine anatomy at the end. The penalty follows as :math:`\rho_k = \lambda/\sigma_k^2`, which keeps the balance between data consistency and denoising that :math:`\lambda` sets. Prerequisites: lesson 10, :doc:`../06-learning/01-plug-and-play`, and lesson 11, :doc:`../06-learning/02-modl-with-admm`. **Learning objectives** - Train a denoiser conditioned on the noise level, ``noise=True`` in :class:`bartorch.learning.UNet`, with :class:`bartorch.learning.Reconstruction`. - Give :class:`bartorch.priors.ImplicitPrior` a schedule of noise levels and :class:`bartorch.optim.ADMMBlock` the matching schedule of penalties. - Compare an annealed schedule with a fixed noise level, iteration by iteration, and apply the same denoiser at a higher acceleration. .. GENERATED FROM PYTHON SOURCE LINES 41-189 .. GENERATED FROM PYTHON SOURCE LINES 190-197 Data ---- You use the slices, coils and fourfold undersampling of :doc:`02-staged-training`: BrainWeb [#brainweb]_ subject 0 to train the denoiser on, subject 4 to reconstruct. The acquisition is wrapped in a function, because you will change the acceleration at the end. .. GENERATED FROM PYTHON SOURCE LINES 197-245 .. code-block:: Python import lightning import torch import bartorch import bartorch.tools as bt from bartorch import learning, linop, optim, priors from monai.metrics import PSNRMetric from torch.utils.data import DataLoader COILS = 8 ITERATIONS = 12 NOISE = 0.02 _ = torch.manual_seed(0) train_images = brain_slices(subject=0, count=24) valid_images = brain_slices(subject=4, count=8) coils = bt.coils(t=bt.grid(D=(SIZE, SIZE, 1)), n=COILS)[:, 0] sensitivities = coils / bartorch.rss(coils, axes=(0,), keepdim=True) def acquisition(acceleration, seed): """Cartesian SENSE, one line in ``acceleration`` on average.""" ky = torch.arange(SIZE) - SIZE / 2 density = torch.exp(-0.5 * (ky / (SIZE / 6)) ** 2) chance = density / density.sum() * (SIZE / acceleration) draw = torch.rand(SIZE, generator=torch.Generator().manual_seed(seed)) lines = draw < chance lines[SIZE // 2 - 4 : SIZE // 2 + 4] = True pattern = lines.to(torch.complex64)[:, None].expand(SIZE, SIZE) pattern = pattern.contiguous() return linop.CartesianSense(sensitivities, (SIZE, SIZE), pattern=pattern) def measure(A, generator): made = [] for x in valid_images: shape, kind = A.oshape, torch.complex64 noise = torch.randn(shape, dtype=kind, generator=generator) made.append(A(x) + NOISE * noise) return made A = acquisition(4, seed=1) kspace = measure(A, torch.Generator().manual_seed(3)) .. GENERATED FROM PYTHON SOURCE LINES 246-256 Noise-conditioned denoiser -------------------------- With ``noise=True`` the U-Net takes the noise level as an input and modulates its features with it, so one network covers a range of SNRs. You train it on pairs of a slice and the same slice with complex white Gaussian noise, at a standard deviation drawn log-uniformly between 0.5 and 20 per cent of the image's peak for each pair (:class:`~bartorch.learning.ComplexNet` scales each image to unit peak). No coil sensitivities, sampling pattern or k-space enter the training. .. GENERATED FROM PYTHON SOURCE LINES 256-285 .. code-block:: Python LOW, HIGH = 0.005, 0.2 network = learning.UNet(2, spatial=2, widths=(16, 32, 64), noise=True) denoiser = learning.ComplexNet(network, spatial=2) def pairs(images): made = [] for x in images: sigma = LOW * (HIGH / LOW) ** torch.rand(()) noisy = x + sigma * torch.randn_like(x) made.append({"input": noisy, "target": x, "sigma": sigma.reshape(1)}) return made trainer = lightning.Trainer( max_epochs=30, accelerator="cpu", logger=False, enable_checkpointing=False, enable_model_summary=False, enable_progress_bar=False, ) trainer.fit( learning.Reconstruction(denoiser, "denoiser", lr=2e-3), DataLoader(train_images, batch_size=4, shuffle=True, collate_fn=pairs), DataLoader(valid_images, batch_size=4, collate_fn=pairs), ) .. rst-class:: sphx-glr-script-out .. code-block:: none /opt/hostedtoolcache/Python/3.12.15/x64/lib/python3.12/site-packages/lightning/pytorch/utilities/_pytree.py:21: `isinstance(treespec, LeafSpec)` is deprecated, use `isinstance(treespec, TreeSpec) and treespec.is_leaf()` instead. /opt/hostedtoolcache/Python/3.12.15/x64/lib/python3.12/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:434: The 'val_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=3` in the `DataLoader` to improve performance. /opt/hostedtoolcache/Python/3.12.15/x64/lib/python3.12/site-packages/lightning/pytorch/trainer/connectors/data_connector.py:434: The 'train_dataloader' does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` to `num_workers=3` in the `DataLoader` to improve performance. .. GENERATED FROM PYTHON SOURCE LINES 286-288 You score every image by the mean PSNR of its magnitude over the eight slices of subject 4. First the denoiser on its own, at three noise levels: .. GENERATED FROM PYTHON SOURCE LINES 288-305 .. code-block:: Python psnr = PSNRMetric(max_val=1.0) truth = torch.stack(valid_images).abs()[:, None] def score(images): return float(psnr(images.abs()[:, None], truth).mean()) with torch.no_grad(): for sigma in (0.01, 0.05, 0.1): noisy = torch.stack(valid_images) noisy = noisy + sigma * torch.randn_like(noisy) level = torch.tensor([sigma]) denoised = torch.stack([denoiser(x[None], level)[0] for x in noisy]) before, after = score(noisy), score(denoised) print(f"sigma {sigma}: {before:5.2f} dB noisy, {after:5.2f} denoised") .. rst-class:: sphx-glr-script-out .. code-block:: none sigma 0.01: 41.46 dB noisy, 41.13 denoised sigma 0.05: 27.17 dB noisy, 27.16 denoised sigma 0.1: 21.06 dB noisy, 21.06 denoised .. GENERATED FROM PYTHON SOURCE LINES 306-323 The gain grows with the noise level: about 2 dB at 0.01, about 8 dB at 0.1. One set of weights covers the whole range, which is what lets the noise level change from one iteration to the next below. Annealing schedules ------------------- :class:`~bartorch.priors.ImplicitPrior` takes a sequence of noise levels, one per iteration, and :class:`~bartorch.optim.ADMMBlock` a sequence of penalties; the last value of each is repeated past its end. The annealed schedule decreases the noise level geometrically from 0.1 to 0.01 of the peak over twelve iterations, and the penalty follows as :math:`\rho_k = \lambda/\sigma_k^2`. When the penalty changes, the block rescales ADMM's scaled dual variable by the ratio of the old penalty to the new, so the unscaled one carries over. The two fixed schedules hold the noise level at 0.03 and at 0.01, with the penalty given by the same :math:`\lambda`. .. GENERATED FROM PYTHON SOURCE LINES 323-339 .. code-block:: Python LAMBDA = 1e-4 schedules = { "annealed, 0.1 to 0.01": torch.logspace(-1.0, -2.0, ITERATIONS), "fixed, 0.03": torch.full((ITERATIONS,), 0.03), "fixed, 0.01": torch.full((ITERATIONS,), 0.01), } def plug_and_play(sigma, lam=LAMBDA): prior = priors.ImplicitPrior(denoiser, sigma=sigma.tolist()) rho = (lam / sigma**2).tolist() block = optim.ADMMBlock(prior, rho=rho, cg_maxiter=5) return learning.Unrolled(block, iterations=ITERATIONS) .. GENERATED FROM PYTHON SOURCE LINES 340-342 :meth:`~bartorch.learning.Unrolled.steps` yields every iterate, so you can follow each schedule iteration by iteration: .. GENERATED FROM PYTHON SOURCE LINES 342-366 .. code-block:: Python curves, results = {}, {} with torch.no_grad(): for name, sigma in schedules.items(): runs = [list(plug_and_play(sigma).steps(y, A)) for y in kspace] iterates = [torch.stack(images) for images in zip(*runs)] curves[name] = [score(images) for images in iterates] results[name] = iterates[-1] cg = torch.stack([optim.cg(y, A, maxiter=20) for y in kspace]) results["CG SENSE, 20 iterations"] = cg .. image-sg:: /auto_examples/09-learning-workflows/images/sphx_glr_04-annealed-plug-and-play_001.png :alt: 04 annealed plug and play :srcset: /auto_examples/09-learning-workflows/images/sphx_glr_04-annealed-plug-and-play_001.png :class: sphx-glr-single-img .. rst-class:: sphx-glr-script-out .. code-block:: none annealed, 0.1 to 0.01 PSNR 27.01 dB fixed, 0.03 PSNR 27.14 dB fixed, 0.01 PSNR 27.48 dB CG SENSE, 20 iterations PSNR 24.21 dB .. GENERATED FROM PYTHON SOURCE LINES 367-378 The fixed level of 0.03 converges within about four iterations, to an image limited by the smoothing the denoiser applies at that level. The fixed level of 0.01 preserves more detail, but its large penalty makes each ADMM step move little from the previous one: it starts far below CG-SENSE and is still rising after twelve iterations. The annealed schedule is ahead from the first iteration, and ends slightly ahead of the better fixed level, without a level having been tuned for this acquisition. In the error maps of one slice, the CG-SENSE error is spread over the whole field of view; the plug-and-play error is confined to the head and largest at the cortex. .. GENERATED FROM PYTHON SOURCE LINES 378-391 .. rst-class:: sphx-glr-horizontal * .. image-sg:: /auto_examples/09-learning-workflows/images/sphx_glr_04-annealed-plug-and-play_002.png :alt: reference, CG-SENSE, annealed plug-and-play :srcset: /auto_examples/09-learning-workflows/images/sphx_glr_04-annealed-plug-and-play_002.png :class: sphx-glr-multi-img * .. image-sg:: /auto_examples/09-learning-workflows/images/sphx_glr_04-annealed-plug-and-play_003.png :alt: CG-SENSE NRMSE 0.124, annealed plug-and-play NRMSE 0.090 :srcset: /auto_examples/09-learning-workflows/images/sphx_glr_04-annealed-plug-and-play_003.png :class: sphx-glr-multi-img .. GENERATED FROM PYTHON SOURCE LINES 392-397 Transfer to another acquisition ------------------------------- You trained the denoiser without an encoding, so it applies unchanged to :math:`R = 6`, a sampling pattern it has never been used with: .. GENERATED FROM PYTHON SOURCE LINES 397-415 .. code-block:: Python A6 = acquisition(6, seed=2) kspace6 = measure(A6, torch.Generator().manual_seed(4)) annealed_pnp = plug_and_play(schedules["annealed, 0.1 to 0.01"]) with torch.no_grad(): annealed = torch.stack([annealed_pnp(y, A6) for y in kspace6]) cg6 = torch.stack([optim.cg(y, A6, maxiter=20) for y in kspace6]) print(f"R = 6: annealed plug-and-play {score(annealed):5.2f} dB") print(f"R = 6: CG SENSE {score(cg6):5.2f} dB") .. rst-class:: sphx-glr-horizontal * .. image-sg:: /auto_examples/09-learning-workflows/images/sphx_glr_04-annealed-plug-and-play_004.png :alt: reference, CG-SENSE, R = 6, annealed, R = 6 :srcset: /auto_examples/09-learning-workflows/images/sphx_glr_04-annealed-plug-and-play_004.png :class: sphx-glr-multi-img * .. image-sg:: /auto_examples/09-learning-workflows/images/sphx_glr_04-annealed-plug-and-play_005.png :alt: CG-SENSE, R = 6 NRMSE 0.168, annealed, R = 6 NRMSE 0.139 :srcset: /auto_examples/09-learning-workflows/images/sphx_glr_04-annealed-plug-and-play_005.png :class: sphx-glr-multi-img .. rst-class:: sphx-glr-script-out .. code-block:: none R = 6: annealed plug-and-play 23.72 dB R = 6: CG SENSE 21.98 dB .. GENERATED FROM PYTHON SOURCE LINES 416-429 At this acceleration eight coils no longer unfold the aliasing well: CG-SENSE loses about 2 dB with respect to :math:`R = 4`, and its error map shows the folded outline of the head on top of the g-factor noise. The plug-and-play reconstruction loses a fraction of a decibel and keeps the anatomy. Exercise: the weight :math:`\lambda` ------------------------------------ :math:`\lambda` sets the weight of the denoising step relative to data consistency at every noise level. Rerun the annealed schedule at :math:`R = 4` with a tenth and ten times the weight above; the denoiser stays as it is: .. GENERATED FROM PYTHON SOURCE LINES 429-435 .. code-block:: Python for lam in (LAMBDA / 10, LAMBDA, LAMBDA * 10): pnp = plug_and_play(schedules["annealed, 0.1 to 0.01"], lam=lam) with torch.no_grad(): made = torch.stack([pnp(y, A) for y in kspace]) print(f"lambda {lam:.0e}: PSNR {score(made):5.2f} dB") .. rst-class:: sphx-glr-script-out .. code-block:: none lambda 1e-05: PSNR 27.01 dB lambda 1e-04: PSNR 27.01 dB lambda 1e-03: PSNR 27.37 dB .. GENERATED FROM PYTHON SOURCE LINES 436-458 Both a tenth and ten times the weight score more than 1 dB below :math:`\lambda = 10^{-4}`. A smaller :math:`\lambda` lowers every penalty :math:`\rho_k`, so the denoiser's output enters the data-consistency step with less weight; a larger one weights it more, and the reconstruction moves away from the measured data. As a spec --------- What this tour built, stated the way you would ask an agent for it: .. code-block:: text Train a noise-conditioned denoiser, ComplexNet(UNet(2, spatial=2, widths=(16, 32, 64), noise=True)), on clean complex brain slices with complex Gaussian noise of a log-uniform level in [0.005, 0.2], 30 epochs of learning.Reconstruction(..., "denoiser", lr=2e-3). Use it as the proximal step of 12 ADMM iterations (optim.ADMMBlock, 5 CG steps) for 2D Cartesian SENSE at R = 4: priors.ImplicitPrior with sigma decreasing geometrically from 0.1 to 0.01 and rho_k = 1e-4 / sigma_k^2. Compare it with fixed sigma 0.03 and 0.01 and with CG-SENSE by PSNR per iteration, and rerun it unchanged at R = 6. .. GENERATED FROM PYTHON SOURCE LINES 460-476 References ---------- .. [#pnp] Venkatakrishnan SV, Bouman CA, Wohlberg B. Plug-and-play priors for model based reconstruction. *IEEE Global Conference on Signal and Information Processing* 945-948 (2013). https://doi.org/10.1109/GlobalSIP.2013.6737048 .. [#dpir] Zhang K, Li Y, Zuo W, Zhang L, Van Gool L, Timofte R. Plug-and-play image restoration with deep denoiser prior. *IEEE Trans Pattern Anal Mach Intell* 44(10):6360-6376 (2022). https://doi.org/10.1109/TPAMI.2021.3088914 .. [#brainweb] Collins DL, Zijdenbos AP, Kollokian V, Sled JG, Kabani NJ, Holmes CJ, Evans AC. Design and construction of a realistic digital brain phantom. *IEEE Trans Med Imaging* 17(3):463-468 (1998). https://doi.org/10.1109/42.712135 .. rst-class:: sphx-glr-timing **Total running time of the script:** (0 minutes 43.403 seconds) .. _sphx_glr_download_auto_examples_09-learning-workflows_04-annealed-plug-and-play.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: 04-annealed-plug-and-play.ipynb <04-annealed-plug-and-play.ipynb>` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: 04-annealed-plug-and-play.py <04-annealed-plug-and-play.py>` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: 04-annealed-plug-and-play.zip <04-annealed-plug-and-play.zip>` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_