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: https://github.com/duong-db/U-Lite.
翻译:卷积神经网络(CNN)和基于Transformer的模型因其提取高层特征和捕获图像重要特征的能力,正被广泛应用于医学图像分割。然而,高精度需求与低计算成本之间往往存在权衡。具有更多参数的模型理论上可取得更优性能,但也会增加计算复杂度和内存占用,因此实际部署面临困难。本文旨在寻找一种保持同等甚至更优性能的轻量级基于U-Net的模型,即U-Lite。我们基于深度可分离卷积原理设计U-Lite,使其既能发挥CNN优势又能显著减少计算参数。具体而言,我们在编码器和解码器中引入7×7核的轴向深度卷积以扩大模型感受野。为进一步提升性能,我们在瓶颈层采用多个3×3滤波器的轴向空洞深度卷积作为分支之一。整体而言,U-Lite仅包含87.8万参数,比传统U-Net减少35倍,更比现代其他基于Transformer的模型减少数十倍参数。该模型在显著降低计算复杂度的同时,在医学分割任务上取得了与当前最优架构相媲美的性能。代码将开源在:https://github.com/duong-db/U-Lite。