High Resolution Medical Image Analysis with Spatial Partitioning
This addresses a bottleneck for medical imaging researchers by enabling end-to-end training on very high-resolution images without information loss, though it is incremental as it builds on existing frameworks like Mesh-TensorFlow.
The paper tackles the problem of training CNN models on high-resolution medical images like 3D CT scans, which are too large for GPU memory, by implementing spatial partitioning to distribute computations across GPUs, enabling training on up to 512x512x512 resolution data.
Medical images such as 3D computerized tomography (CT) scans and pathology images, have hundreds of millions or billions of voxels/pixels. It is infeasible to train CNN models directly on such high resolution images, because neural activations of a single image do not fit in the memory of a single GPU/TPU, and naive data and model parallelism approaches do not work. Existing image analysis approaches alleviate this problem by cropping or down-sampling input images, which leads to complicated implementation and sub-optimal performance due to information loss. In this paper, we implement spatial partitioning, which internally distributes the input and output of convolutional layers across GPUs/TPUs. Our implementation is based on the Mesh-TensorFlow framework and the computation distribution is transparent to end users. With this technique, we train a 3D Unet on up to 512 by 512 by 512 resolution data. To the best of our knowledge, this is the first work for handling such high resolution images end-to-end.