learning.Reconstruction

learning.Reconstruction#

class bartorch.learning.Reconstruction#

Bases: LightningModule

Lightning module training one of the three stages of an unrolled reconstruction network.

  • "denoiser" trains the network alone, on pairs of a degraded image and its reference.

  • "greedy" trains an Unrolled built with detach=True, back-propagating a loss on each iteration’s image before the next iteration runs; the losses grow geometrically, the last ratio times the first.

  • "end-to-end" takes the loss on the last image through the whole stack, which checkpoint=True makes affordable.

Run in this order, each from the weights the previous one left, these are the staged training of Urman et al.

Items are dictionaries and a batch is a list of them (collate_fn=list), since operators and k-space vary between subjects:

  • "denoiser": input and target images, and optionally step, sigma and label passed to the network;

  • the unrolled stages: the data y and its encoding A; the reference image target, or the sampling pattern for self-supervised training; optionally x0.

Any item may carry weights, multiplying the voxel-wise error of the default loss. Batches are not moved to the trainer’s device: a Patchwise network or an operator with operands on the host moves what it needs.

An item without target is trained self-supervised (SSDU): its pattern is partitioned by split() at every step, the network reconstructs from one part, and the loss is the normalized \(\ell_1 + \ell_2\) error of the held-out k-space. The validation partition is the same at every epoch.

Optimization is manual: gradients are accumulated over accumulate items, optionally clipped, and applied with Adam; the learning rate is multiplied by factor after patience epochs without improvement of the logged val_loss.

Parameters:

  • model (nn.Module) – The network for "denoiser", an Unrolled for the other stages.

  • stage ({"denoiser", "greedy", "end-to-end"}, default="end-to-end") – Training stage.

  • loss (callable, default=None) – loss(prediction, target, item) for supervised items; the mean modulus of the error when omitted.

  • ratio (float, default=10.0) – Weight of the last iteration’s loss relative to the first’s, in the greedy stage.

  • lr (float, default=1e-3) – Adam’s learning rate.

  • weight_decay (float, default=0.0) – Adam’s weight decay.

  • factor (float, default=0.5) – Multiplier of the learning rate on a plateau of val_loss.

  • patience (int, default=10) – Epochs without improvement of val_loss before the learning rate is reduced.

  • accumulate (int, default=1) – Items whose gradients are summed before an optimizer step.

  • clip (float, default=None) – Maximum norm of the gradient; no clipping when omitted.

  • fraction (float, default=0.4) – Share of the acquired samples held out, for self-supervised items.

  • split_options (dict, default=None) – Further keywords of split().

Raises:

ValueError – An unknown stage; "greedy" with a model not built with detach=True; accumulate below one. In training, "greedy" with a block whose image does not depend on that iteration’s parameters, such as ADMMBlock, whose image is the conjugate-gradient update preceding the denoiser; ISTBlock qualifies.

References

Urman Y, Nishimura M, Abraham D, 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

Examples

>>> stage = learning.Reconstruction(model, stage="greedy", accumulate=4)
>>> trainer = lightning.Trainer(max_epochs=50, callbacks=[EarlyStopping("val_loss")])
>>> trainer.fit(stage, DataLoader(train, collate_fn=list), DataLoader(valid, collate_fn=list))
transfer_batch_to_device()#

Leave the batch where the dataset put it.

configure_optimizers()#

Adam over the trainable parameters, reduced on a plateau of val_loss.

training_step()#

Back-propagate each item of a list and step every accumulate items.

validation_step()#

Log the mean loss of the items as val_loss.

on_validation_epoch_end()#

Step the plateau scheduler on val_loss.