Mixed-Precision Quantization of U-Net for Medical Image Segmentation

Network quantization can facilitate the practical deployment of U-Net in various medical applications, especially medical image segmentation. However, existing quantization methods for U-Net do not sufficiently exploit the skip connection between its encoder and decoder structures. In this paper, we propose a novel mixed-precision quantization scheme for U-Net that leverages split convolution and skip-supervised quantization aware training to address this problem. Specifically, split convolution decouples the feature maps obtained by skip connection and up-sampling to enable fine-grained bitwidth allocation for mixed-precision quantization. Furthermore, we introduce supervision on skip connection for quantization aware training to compensate gradients in back propagation and enhance the skip features in finetuning. To our best knowledge, this is the first attempt to design mixed-precision quantization specifically for medical image segmentation with U-Net. Experimental results on four medical image datasets demonstrate that the proposed scheme outperforms existing U-Net quantization strategy and recent general-purpose mixed-precision quantizaton methods in medical image segementation.

I. Introduction

With the help of deep learning, automatic medical image segmentation achieves promising results and liberates pathologists from time-consuming and tedious labeling. U-Net [1], a fully convolutional neural network architecture, is widely used for medical image segmentation. The increasing model complexity [2] –[5] for U-Net improvement and diverse deployment requirements raise the demands for U-Net compression.