Low-Memory CNNs Enabling Real-Time Ultrasound Segmentation Towards Mobile Deployment

Convolutional Neural Networks (CNNs), which are currently state-of-the-art for most image analysis tasks, are ill suited to leveraging the key benefits of ultrasound imaging - specifically, ultrasound's portability and real-time capabilities. CNNs have large memory footprints, which obstructs their implementation on mobile devices, and require numerous floating point operations, which results in slow CPU inference times. In this article, we propose three approaches to training efficient CNNs that can operate in real-time on a CPU (catering to a clinical setting), with a low memory footprint, for minimal compromise in accuracy. We first demonstrate the power of 'thin' CNNs (with very few feature channels) for fast medical image segmentation. We then leverage separable convolutions to further speed up inference, reduce parameter count and facilitate mobile deployment. Lastly, we propose a novel knowledge distillation technique to boost the accuracy of light-weight models, while maintaining inference speed-up. For a negligible sacrifice in test set Dice performance on the challenging ultrasound analysis task of nerve segmentation, our final proposed model processes images at 30 fps on a CPU, which is 9x faster than the standard U-Net, while requiring 420x less space in memory. Code for this work is available at: https://github.com/sagarvaze96/lightweight_unet.