1M parameters are enough? A lightweight CNN-based model for medical image segmentation

Binh-Duong Dinh,Thanh-Thu Nguyen,Thi-Thao Tran,Van-Truong Pham

2023-07-03

Abstract:Convolutional neural networks (CNNs) and Transformer-based models are being widely applied in medical image segmentation thanks to their ability to extract high-level features and capture important aspects of the image. However, there is often a trade-off between the need for high accuracy and the desire for low computational cost. A model with higher parameters can theoretically achieve better performance but also result in more computational complexity and higher memory usage, and thus is not practical to implement. In this paper, we look for a lightweight U-Net-based model which can remain the same or even achieve better performance, namely U-Lite. We design U-Lite based on the principle of Depthwise Separable Convolution so that the model can both leverage the strength of CNNs and reduce a remarkable number of computing parameters. Specifically, we propose Axial Depthwise Convolutions with kernels 7x7 in both the encoder and decoder to enlarge the model receptive field. To further improve the performance, we use several Axial Dilated Depthwise Convolutions with filters 3x3 for the bottleneck as one of our branches. Overall, U-Lite contains only 878K parameters, 35 times less than the traditional U-Net, and much more times less than other modern Transformer-based models. The proposed model cuts down a large amount of computational complexity while attaining an impressive performance on medical segmentation tasks compared to other state-of-the-art architectures. The code will be available at: <a class="link-external link-https" href="https://github.com/duong-db/U-Lite" rel="external noopener nofollow">this https URL</a>.

Image and Video Processing,Computer Vision and Pattern Recognition

What problem does this paper attempt to address?

The paper aims to address the demand for lightweight models in medical image segmentation tasks. Specifically, the authors propose a lightweight model based on the U-Net architecture, named U-Lite, which aims to reduce the number of parameters, thereby lowering computational complexity and memory usage while maintaining or even enhancing segmentation performance. Traditionally, high-accuracy models often come with a large number of parameters, leading to high computational costs. U-Lite, by introducing the Axial Depthwise Convolution module, significantly reduces the number of parameters while ensuring high performance. Compared to the traditional U-Net model, U-Lite reduces the number of parameters by 35 times, and experimental results on multiple datasets show that its performance surpasses that of various existing advanced models. Therefore, this study addresses the limitations of existing complex models in practical applications and provides an efficient and lightweight solution.