.. DO NOT EDIT. .. THIS FILE WAS AUTOMATICALLY GENERATED BY SPHINX-GALLERY. .. TO MAKE CHANGES, EDIT THE SOURCE PYTHON FILE: .. "auto_examples/06-learning/02-modl-with-admm.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_06-learning_02-modl-with-admm.py: ================================ 11. Training an unrolled network ================================ In lesson 10 the proximal step of ADMM was a denoiser pretrained on photographs, never trained on MR images or on the residual aliasing of an undersampled scan. In this lesson you train the denoiser through the iteration: a fixed number of ADMM steps is unrolled into one network, which is fitted end to end against fully sampled images. This is MoDL [#modl]_. The encoding model, the SENSE operator :math:`A` with its coil sensitivities, Fourier transform and sampling pattern, stays fixed, and the network is trained only for the part of the image that the data leave undetermined. As published, MoDL alternates a learned denoiser :math:`D_w` and a data-consistency step (half-quadratic splitting): .. math:: z^{k} &= D_w(x^{k}) \\ x^{k+1} &= \arg\min_x \; \|A x - y\|_2^2 + \lambda \|x - z^{k}\|_2^2. The second line is a conjugate-gradient solve of :math:`(A^H A + \lambda) x = A^H y + \lambda z^{k}`. ADMM [#boyd]_ is the same splitting with a dual variable :math:`u` carried along, .. math:: x^{k+1} &= \arg\min_x \; \|A x - y\|_2^2 + \rho \|x - z^{k} + u^{k}\|_2^2 \\ z^{k+1} &= D_w(x^{k+1} + u^{k}) \\ u^{k+1} &= u^{k} + x^{k+1} - z^{k+1}, so MoDL is ADMM with :math:`u` held at zero. The dual variable accumulates the mismatch between the data-consistent and the denoised image, so that a fixed point satisfies the data constraint rather than a penalized version of it. :class:`bartorch.optim.ADMMBlock` is one step of BART's ``admm.c``, conjugate-gradient x-update included, and :class:`bartorch.learning.Unrolled` repeats it. The network enters the proximal step through :class:`bartorch.priors.ImplicitPrior`, as the pretrained one did in lesson 10, and :math:`\rho` is trained with its weights. **Learning objectives** - Put a network in the proximal step of :class:`bartorch.optim.ADMMBlock` through :class:`bartorch.priors.ImplicitPrior`. - Unroll the iteration with :class:`bartorch.learning.Unrolled` and train it end to end with ``lightning``, on data held and augmented by ``torchio``. - Cut the memory of a deeper unrolled network with checkpointing or per-iteration training. Previous: :doc:`01-plug-and-play`. Next: the :doc:`Tours `, where you continue with standalone topics; the learning workflows among them cover networks for complex volumes, staged and self-supervised training, annealed plug-and-play and uncertainty. .. GENERATED FROM PYTHON SOURCE LINES 62-178 .. GENERATED FROM PYTHON SOURCE LINES 179-192 Training images --------------- You need fully sampled images to train against. Here they are 32 axial slices of one BrainWeb subject [#brainweb]_, rendered as :math:`T_1`-weighted spin-echo images with a smooth phase, as in lesson 2; the code is hidden. You split them by position, not at random: the first 24 train and the last 8 validate, so that no validation slice has a near-duplicate neighbour in the training set. One subject, 24 slices and one sampling pattern are a phantom, not a training set. The trained weights are not expected to generalise; the lesson is about the construction. .. GENERATED FROM PYTHON SOURCE LINES 192-209 .. code-block:: Python import lightning import torch import torchio import bartorch import bartorch.tools as bt from bartorch import learning, linop, optim, priors from monai.metrics import PSNRMetric, SSIMMetric from torch.utils.data import DataLoader _ = torch.manual_seed(0) train_images = images[:24] valid_images = images[24:] print(f"{len(train_images)} slices to train, {len(valid_images)} to validate") .. rst-class:: sphx-glr-script-out .. code-block:: none 24 slices to train, 8 to validate .. GENERATED FROM PYTHON SOURCE LINES 210-218 The acquisition --------------- Eight channels of BART's analytical head coil and a variable-density random undersampling of the phase encodes, with the central 8 always acquired. The pattern is the same for every slice, so one operator serves the whole dataset. If each item had its own pattern, you would need one operator per item. .. GENERATED FROM PYTHON SOURCE LINES 218-238 .. code-block:: Python COILS = 8 ACCELERATION = 4 CENTRE = 8 # phase encodes always acquired sensitivities = bt.coils(t=bt.grid(D=(SIZE, SIZE, 1)), n=COILS)[:, 0] sensitivities /= bartorch.rss(sensitivities, axes=(0,), keepdim=True) pattern = lines.to(torch.complex64)[:, None].expand(SIZE, SIZE).contiguous() A = linop.CartesianSense(sensitivities, (SIZE, SIZE), pattern=pattern) print(f"{int(lines.sum())} of {SIZE} phase encodes") print(f"A: {A.ishape} -> {A.oshape}") .. rst-class:: sphx-glr-script-out .. code-block:: none 34 of 128 phase encodes A: (128, 128) -> (8, 128, 128) .. GENERATED FROM PYTHON SOURCE LINES 239-244 The k-space of a slice is :math:`A x` plus complex Gaussian noise. ``A`` is built for one image, so you apply it to a batch item by item. The blocks in :mod:`bartorch.optim` do the same inside, which is why the network below accepts a batch even though the operator does not. ``measure`` also returns :math:`A^H y`, the image the unrolled network starts from. .. GENERATED FROM PYTHON SOURCE LINES 244-269 .. code-block:: Python NOISE = 0.005 def measure(images, generator=None): """Images, their noisy k-space and the adjoint reconstructions.""" x = torch.stack(list(images)) y = torch.stack([A(item) for item in x]) noise = torch.randn(y.shape, dtype=torch.complex64, generator=generator) y = y + NOISE * noise return x, y, torch.stack([A.H(item) for item in y]) .. image-sg:: /auto_examples/06-learning/images/sphx_glr_02-modl-with-admm_001.png :alt: training, training, sampling pattern, validation, $A^H y$ :srcset: /auto_examples/06-learning/images/sphx_glr_02-modl-with-admm_001.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 270-289 Two training slices, the sampling pattern (phase encoding vertical), and what the network starts from on a validation slice: :math:`A^H y`, with the aliasing of the random undersampling and the noise. The dataset ----------- ``torchio`` holds the images and augments them. Its ``ScalarImage`` wants a real tensor of shape ``(channels, width, height, depth)``; :func:`bartorch.learning.as_real` and :func:`~bartorch.learning.as_complex` convert between that and a complex image, with the real and imaginary parts as the two channels and the slice as a volume one voxel deep. Why ``torchio`` and not a plain list? It draws one random transform per subject and applies it to every image of that subject, so anything that has to stay registered with the image (coil sensitivities, parameter maps) is transformed with it. Here each subject holds one image, and the transform is a flip and a rotation of at most eight degrees: the tissue the denoiser learns stays the same, its orientation varies. .. GENERATED FROM PYTHON SOURCE LINES 289-326 .. code-block:: Python augmentation = torchio.Compose( [ torchio.RandomFlip(axes=(0,), flip_probability=0.5), torchio.RandomAffine( scales=0, degrees=(0, 0, 0, 0, -8, 8), translation=0, ), ] ) def subjects(images): return [ torchio.Subject(image=torchio.ScalarImage(tensor=real)) for real in (learning.as_real(image)[..., None] for image in images) ] def collate(batch): """Simulate k-space for a batch of subjects.""" planes = (subject["image"][torchio.DATA][..., 0] for subject in batch) return measure(learning.as_complex(plane) for plane in planes) train_loader = DataLoader( torchio.SubjectsDataset(subjects(train_images), transform=augmentation), batch_size=2, shuffle=True, collate_fn=collate, ) valid_loader = DataLoader( torchio.SubjectsDataset(subjects(valid_images)), batch_size=2, collate_fn=collate, ) .. GENERATED FROM PYTHON SOURCE LINES 327-346 The network ----------- Three objects: * ``deepinv``'s ``DnCNN`` [#dncnn]_, a residual convolutional denoiser of the family MoDL's own five-layer network belongs to. It is a plain ``nn.Module`` on real images. * :class:`~bartorch.priors.ImplicitPrior` presents it as a regularization term. ``spatial=2, channels=2`` sets the layout: real and imaginary parts as two channels, the batch axes folded, and each image scaled to unit peak modulus around the call. * :class:`~bartorch.learning.Unrolled` applies the ADMM step ``ITERATIONS`` times. One block is shared by every iteration, which is MoDL's weight sharing, and ``rho``, MoDL's :math:`\lambda`, is made trainable. ``alpha=1.0`` turns off BART's over-relaxation, so the step is the iteration written at the top; ``cg_maxiter`` is the budget of the x-update, MoDL's ten conjugate-gradient steps. .. GENERATED FROM PYTHON SOURCE LINES 346-365 .. code-block:: Python from deepinv.models import DnCNN ITERATIONS = 5 # unrolled steps, MoDL's K def modl(): """An unrolled network, and the block it shares.""" network = DnCNN(in_channels=2, out_channels=2, depth=5, pretrained=None) prior = priors.ImplicitPrior(network, spatial=2, channels=2) block = optim.ADMMBlock(prior, rho=0.05, alpha=1.0, cg_maxiter=10) block.rho.requires_grad_() return learning.Unrolled(block, iterations=ITERATIONS), block model, block = modl() learned = sum(p.numel() for p in model.parameters() if p.requires_grad) print(f"{learned} learned values, shared by all {ITERATIONS} iterations") print(f"rho starts at {float(block.rho.detach()):.3f}") .. rst-class:: sphx-glr-script-out .. code-block:: none 113155 learned values, shared by all 5 iterations rho starts at 0.050 .. GENERATED FROM PYTHON SOURCE LINES 366-374 Training -------- ``lightning`` runs the loop, and the module is an ordinary supervised one: a forward pass, a loss against the fully sampled image, and metrics from ``monai``. The loss is computed on a tensor and requires nothing of the reconstruction that produced it. The gradient is propagated through the denoiser, the conjugate-gradient solves and :math:`\rho`. .. GENERATED FROM PYTHON SOURCE LINES 374-419 .. code-block:: Python EPOCHS = 15 psnr = PSNRMetric(max_val=1.0) ssim = SSIMMetric(spatial_dims=2, data_range=1.0) class Reconstruction(lightning.LightningModule): """Supervised training of an unrolled network.""" def __init__(self, model, lr=1e-3): super().__init__() self.model = model self.lr = lr def forward(self, y, start): return self.model(y, A, x0=start) def training_step(self, batch, index): x, y, start = batch loss = (self(y, start) - x).abs().square().mean() self.log("loss", loss, prog_bar=True) return loss def validation_step(self, batch, index): x, y, start = batch out = self(y, start).abs()[:, None] self.log("psnr", psnr(out, x.abs()[:, None]).mean(), prog_bar=True) self.log("ssim", ssim(out, x.abs()[:, None]).mean(), prog_bar=True) def configure_optimizers(self): return torch.optim.Adam(self.parameters(), lr=self.lr) trainer = lightning.Trainer( max_epochs=EPOCHS, accelerator="cpu", logger=False, enable_checkpointing=False, enable_model_summary=False, gradient_clip_val=1.0, ) trainer.fit(Reconstruction(model), train_loader, valid_loader) print(f"rho ended at {float(block.rho.detach()):.3f}") .. 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. Epoch 14/14 ━━━━━━━━━━━━━━━━━ 12/12 0:00:16 • 0:00:00 0.74it/s loss: 0.000 psnr: 32.696 ssim: 0.885 rho ended at 0.027 .. GENERATED FROM PYTHON SOURCE LINES 420-427 Results ------- You compare the trained network with three reconstructions of the same validation k-space: the adjoint it starts from; conjugate-gradient SENSE, which fits the data term alone; and the same ADMM iteration with a wavelet penalty in place of the network, run for fifty iterations rather than five. .. GENERATED FROM PYTHON SOURCE LINES 427-477 .. code-block:: Python torch.manual_seed(7) truth, kspace, adjoint = measure(valid_images) with torch.no_grad(): unrolled = model(kspace, A, x0=adjoint) cg = torch.stack([optim.cg(item, A, maxiter=30) for item in kspace]) wavelet = priors.Wavelet(axes=(-1, -2), weight=0.002) admm = torch.stack( [optim.admm(item, A, wavelet, maxiter=50, rho=0.1) for item in kspace], ) def quality(estimate): """Mean PSNR and SSIM of a batch of magnitudes against the truth.""" a, b = estimate.abs()[:, None], truth.abs()[:, None] return float(psnr(a, b).mean()), float(ssim(a, b).mean()) rows = { "adjoint": adjoint, "CG SENSE": cg, "ADMM, wavelet": admm, f"MoDL, K={ITERATIONS}": unrolled, } for name, estimate in rows.items(): made = quality(estimate) print(f"{name:>16} PSNR {made[0]:5.2f} dB SSIM {made[1]:.3f}") .. rst-class:: sphx-glr-horizontal * .. image-sg:: /auto_examples/06-learning/images/sphx_glr_02-modl-with-admm_002.png :alt: CG SENSE error, ADMM, wavelet error, MoDL, K=5 error :srcset: /auto_examples/06-learning/images/sphx_glr_02-modl-with-admm_002.png :class: sphx-glr-multi-img * .. image-sg:: /auto_examples/06-learning/images/sphx_glr_02-modl-with-admm_003.png :alt: truth, CG SENSE, ADMM, wavelet, MoDL, K=5 :srcset: /auto_examples/06-learning/images/sphx_glr_02-modl-with-admm_003.png :class: sphx-glr-multi-img .. rst-class:: sphx-glr-script-out .. code-block:: none adjoint PSNR 23.71 dB SSIM 0.679 CG SENSE PSNR 29.69 dB SSIM 0.681 ADMM, wavelet PSNR 31.64 dB SSIM 0.898 MoDL, K=5 PSNR 32.70 dB SSIM 0.884 .. GENERATED FROM PYTHON SOURCE LINES 478-527 The error maps share one scale. CG SENSE removes most of the aliasing but amplifies the noise, which covers the whole field of view. The wavelet penalty and the unrolled network both suppress it, and both leave their largest errors at the bright, thin scalp. In the enlarged region the two differ little; the network has the higher PSNR, the wavelet penalty the slightly higher SSIM. Training also lowered :math:`\rho` from 0.05 to about 0.026, which weights the network's output less against the data in each x-update than at the start. Do not read the table as a comparison of methods. Fifteen epochs over 24 slices of one subject, against a wavelet penalty with a hand-picked weight, says nothing about either on measured data. For that you would need many subjects, validation on subjects left out of training, and a fixed reconstruction time. Deeper unrolled stack --------------------- Five iterations of a 2D encoding record a computational graph of modest size. Ten iterations of a 3D non-Cartesian encoding do not, and the memory is dominated by the denoiser. The backward pass of an operator is one more application of that operator and stores nothing that grows with the iteration count; a convolutional network stores every activation, once per iteration. :class:`~bartorch.learning.Unrolled` provides two alternatives, neither of which changes the forward value (:doc:`../../explanation/learned-reconstruction`): * ``detach=True`` starts each iteration from a detached state, so the graph spans one iteration. With a loss on each image that :meth:`~bartorch.learning.Unrolled.steps` yields, that is greedy per-iteration training, and its memory does not depend on the iteration count. * ``checkpoint=True`` keeps only the states between iterations and recomputes the interior of each step during the backward pass. The gradient is the end-to-end one, and each block is applied twice. Pretraining the denoiser alone, then greedy per-iteration training, then end-to-end fine-tuning with checkpointing is the staged schedule reported for a fully 3D unrolled reconstruction [#urman]_; the Tour :doc:`../09-learning-workflows/02-staged-training` follows it. Greedy training does not suit the ADMM step: its image is the x-update, which depends on the denoiser only through the previous iteration's auxiliary variable, and a detached start cuts that dependence. That Tour uses a proximal-gradient step instead, whose image is the denoiser's output. Compare the gradient of ``rho`` with and without checkpointing. ``rho`` enters every iteration and every conjugate-gradient solve, so its gradient is propagated through all of them: .. GENERATED FROM PYTHON SOURCE LINES 527-543 .. code-block:: Python x, y, start = measure(valid_images[:1]) made = [] for recompute in (False, True): stack = learning.Unrolled( block, iterations=ITERATIONS, checkpoint=recompute, ) for parameter in stack.parameters(): parameter.grad = None (stack(y, A, x0=start) - x).abs().square().mean().backward() made.append(float(block.rho.grad)) print(f"rho's gradient: {made[0]:.6g} recorded, {made[1]:.6g} recomputed") .. rst-class:: sphx-glr-script-out .. code-block:: none rho's gradient: 0.000585573 recorded, 0.000585573 recomputed .. GENERATED FROM PYTHON SOURCE LINES 544-559 The two gradients agree to every printed digit: checkpointing changes the memory, not the result. A third alternative is not to unroll. :class:`bartorch.optim.FixedPoint` runs the block to its fixed point and differentiates there, by solving the adjoint fixed-point equation, so its memory is that of one step whatever the iteration count. That is a deep-equilibrium model [#deq]_, and the stack you trained is its truncated form. Exercise: iterations beyond training ------------------------------------ The block was trained for five iterations. Wrap the same trained block in a new :class:`~bartorch.learning.Unrolled` with fewer or more iterations and evaluate it on the validation slices; this needs forward passes only. .. GENERATED FROM PYTHON SOURCE LINES 559-578 .. code-block:: Python depths = (1, 2, 3, 5, 8, 12) scores = [] with torch.no_grad(): for depth in depths: stack = learning.Unrolled(block, iterations=depth) scores.append(quality(stack(kspace, A, x0=adjoint))[0]) .. image-sg:: /auto_examples/06-learning/images/sphx_glr_02-modl-with-admm_004.png :alt: 02 modl with admm :srcset: /auto_examples/06-learning/images/sphx_glr_02-modl-with-admm_004.png :class: sphx-glr-single-img .. GENERATED FROM PYTHON SOURCE LINES 579-606 PSNR rises up to the five iterations the block was trained for and falls beyond them: at 8 and 12 iterations it is below the wavelet penalty again. Training fitted the block to the iterates of a five-step stack and constrained nothing after the fifth, so the unrolled network is not a convergent iteration. If the iteration count has to be free at inference, :class:`~bartorch.optim.FixedPoint` is the formulation that trains for it. As a spec --------- What this lesson built, stated the way you would ask an agent for it: .. code-block:: text Train a MoDL-style unrolled reconstruction for 2D Cartesian SENSE (128x128, 8 coils, variable-density pattern of 34 of 128 lines including the central 8, known sensitivities, k-space noise 0.005) with bartorch. Wrap deepinv's DnCNN (2 in/out channels, depth 5, untrained) as priors.ImplicitPrior(spatial=2, channels=2), put it in optim.ADMMBlock(rho=0.05, alpha=1.0, cg_maxiter=10) with rho trainable, and unroll 5 iterations with learning.Unrolled. Hold the images in torchio (random flip, rotation up to 8 degrees), simulate k-space in the collate function, and train with lightning (Adam, lr 1e-3, 15 epochs, batch 2, squared error against the fully sampled image, monai PSNR and SSIM on validation). Compare against CG SENSE and wavelet ADMM, and check that checkpoint=True reproduces the gradient of rho. .. GENERATED FROM PYTHON SOURCE LINES 608-637 References ---------- .. [#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 .. [#modl] Aggarwal HK, Mani MP, Jacob M. MoDL: model-based deep learning architecture for inverse problems. *IEEE Trans Med Imaging* 38(2):394-405 (2019). https://doi.org/10.1109/TMI.2018.2865356 .. [#boyd] Boyd S, Parikh N, Chu E, Peleato B, Eckstein J. Distributed optimization and statistical learning via the alternating direction method of multipliers. *Found Trends Mach Learn* 3(1):1-122 (2011). https://doi.org/10.1561/2200000016 .. [#dncnn] Zhang K, Zuo W, Chen Y, Meng D, Zhang L. Beyond a Gaussian denoiser: residual learning of deep CNN for image denoising. *IEEE Trans Image Process* 26(7):3142-3155 (2017). https://doi.org/10.1109/TIP.2017.2662206 .. [#urman] Urman Y, Nishimura M, Abraham DR, Cao X, Setsompop K. Fully 3D unrolled magnetic resonance fingerprinting reconstruction via staged pretraining and implicit gridding. *Magn Reson Med* 96(5):2516-2529 (2026). https://doi.org/10.1002/mrm.70500 .. [#deq] Bai S, Kolter JZ, Koltun V. Deep equilibrium models. *Advances in Neural Information Processing Systems* 32:688-699 (2019). https://arxiv.org/abs/1909.01377 .. rst-class:: sphx-glr-timing **Total running time of the script:** (5 minutes 1.130 seconds) .. _sphx_glr_download_auto_examples_06-learning_02-modl-with-admm.py: .. only:: html .. container:: sphx-glr-footer sphx-glr-footer-example .. container:: sphx-glr-download sphx-glr-download-jupyter :download:`Download Jupyter notebook: 02-modl-with-admm.ipynb <02-modl-with-admm.ipynb>` .. container:: sphx-glr-download sphx-glr-download-python :download:`Download Python source code: 02-modl-with-admm.py <02-modl-with-admm.py>` .. container:: sphx-glr-download sphx-glr-download-zip :download:`Download zipped: 02-modl-with-admm.zip <02-modl-with-admm.zip>` .. only:: html .. rst-class:: sphx-glr-signature `Gallery generated by Sphinx-Gallery `_