Computational imaging systems jointly design computation and hardware to retrieve information which is not traditionally accessible with standard imaging systems. Recently, critical aspects such as experimental design and image priors are optimized through deep neural networks formed by the unrolled iterations of classical physics-based reconstructions (termed physics-based networks). However, for real-world large-scale systems, computing gradients via backpropagation restricts learning due to memory limitations of graphical processing units. In this work, we propose a memory-efficient learning procedure that exploits the reversibility of the network’s layers to enable data-driven design for large-scale computational imaging. We demonstrate our methods practicality on two large-scale systems: super-resolution optical microscopy and multi-channel magnetic resonance imaging. Computational imaging systems (tomographic systems, computational optics, magnetic resonance imaging, to name a few) jointly design software and hardware to retrieve information which is not traditionally accessible on standard imaging systems. Generally, such systems are characterized by how the information is encoded (forward process) and decoded (inverse problem) from the measurements. The decoding process is typically iterative in nature, alternating between enforcing data consistency and image prior knowledge. Recent work has demonstrated the ability to optimize computational imaging systems by unrolling the iterative decoding process to form a differentiable Physicsbased Network (PbN) (1; 2; 3) and then relying on a dataset and training to learn the system's design parameters, e.g. experimental design (3; 4; 5) , image prior model (1; 2; 6; 7). PbNs are constructed from the operations of reconstruction, e.g. proximal gradient descent algorithm. By including known structures and quantities, such as the forward model, gradient, and proximal updates, PbNs can be efficiently parameterized by only a few learnable variables, thereby enabling an efficient use of training data (6) while still retaining robustness associated with conventional physics-based inverse problems. Training PbNs relies on gradient-based updates computed using backpropagation (an implementation of reverse-mode differentiation (8) ). Most modern imaging systems seek to decode ever-larger growing quantities of information (gigabytes to terabytes) and as this grows, memory required to perform backpropagation is limited by the memory capacity of modern graphical processing units (GPUs). Methods to save memory during backpropagation (e.g. forward recalculation, reverse recalculation, and checkpointing) trade off spatial and temporal complexity (8) . For a PbN with N layers, standard backpropagation achieves O(N ) temporal and spatial complexity. Forward recalculation achieves O(1) memory complexity, but has to recalculate unstored variables forward from the input of the network when needed, yielding O(N 2 ) temporal complexity. Forward checkpointing smoothly trades off temporal, O(N K), and spatial, O(N/K), complexity by saving variables every K layers and forward-recalculating unstored variables from the closest checkpoint. Reverse recalculation provides a practical solution to beat the trade off between spatial vs. temporal complexity by calculating unstored variables in reverse from the output of the network, yielding O(N ) temporal and O(1) spatial complexities. Recently, several reversibility schemes have been proposed for residual networks (9), learning ordinary differential equations (10) , and other specialized network architectures (11; 12) . In this work, we propose a memory-efficient learning procedure for backpropagation for the PbN formed from proximal gradient descent, thereby enabling learning for many large-scale computational imaging systems. Based on the concept of invertibility and reverse recalculation, we detail how backpropagation can be performed without the need to store intermediate variables for networks composed of gradient and proximal layers. We highlight practical restrictions on the layers and introduce a hybrid scheme that combines our reverse recalculation methods with checkpointing to mitigate numerical error accumulation. Finally, we demonstrate our method's usefulness to learn the design for two practical large-scale computational imaging systems: superresolution optical microscopy (Fourier Ptychography) and multi-channel magnetic resonance imaging. In this communication, we presented a practical memory-efficient learning method for large-scale computational imaging problems without dramatically increasing training time. Using the concept of reversibility, we implemented reverse-mode differentiation with favorable spatial and temporal complexities. We demonstrated our method on two representative applications: SR optical microscopy and multi-channel MRI. We expect other computational imaging systems to nicely fall within our framework.
We propose a memory-efficient learning procedure that exploits the reversibility of the network’s layers to enable data-driven design for large-scale computational imaging.
{"src": "scitldr-AIC_train", "paper_id": "ryeo27n9Ur"}