A Partially Reversible U-Net for Memory-Efficient Volumetric Image Segmentation
This addresses memory bottlenecks in medical imaging segmentation, allowing for more efficient and accurate 3D models, though it is an incremental improvement based on reversible networks.
The paper tackles the high memory consumption of 3D convolutional neural networks for volumetric image segmentation by proposing a partially reversible U-Net architecture, which reduces memory usage substantially and enables deeper networks with higher segmentation accuracy on the BraTS dataset.
One of the key drawbacks of 3D convolutional neural networks for segmentation is their memory footprint, which necessitates compromises in the network architecture in order to fit into a given memory budget. Motivated by the RevNet for image classification, we propose a partially reversible U-Net architecture that reduces memory consumption substantially. The reversible architecture allows us to exactly recover each layer's outputs from the subsequent layer's ones, eliminating the need to store activations for backpropagation. This alleviates the biggest memory bottleneck and enables very deep (theoretically infinitely deep) 3D architectures. On the BraTS challenge dataset, we demonstrate substantial memory savings. We further show that the freed memory can be used for processing the whole field-of-view (FOV) instead of patches. Increasing network depth led to higher segmentation accuracy while growing the memory footprint only by a very small fraction, thanks to the partially reversible architecture.