[论文解读] 1M parameters are enough? A lightweight CNN-based model for medical image segmentation

Binh-Duong Dinh, Thanh-Thu Nguyen+2 more|arXiv (Cornell University)|Jun 28, 2023

Advanced Neural Network Applications被引用 5

一句话总结

该论文提出U-Lite，一种基于U-Net的轻量化卷积神经网络，参数量仅为87.8万，比标准U-Net少35倍，通过使用轴向深度可分离卷积和空洞卷积，在医学图像分割任务中保持了高性能。该模型在基准数据集上实现了最先进性能，同时计算成本显著降低，证明了100万参数足以实现高精度的医学图像分割。

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: https://github.com/duong-db/U-Lite.

研究动机与目标

• 开发一种基于轻量化卷积神经网络的医学图像分割模型，在大幅降低计算成本的同时保持高精度。
• 探究参数量少于100万的模型是否能够匹配或超越更大、更复杂的模型架构的性能。
• 通过利用深度可分离卷积和轴向卷积结构设计，降低模型复杂度而不损失分割精度。
• 构建一种适用于资源受限临床环境的实用且可部署的模型。
• 为标准U-Net和现代基于Transformer的模型提供一种参数高效的替代方案，用于医学图像分割。

提出的方法

• 设计一种类似U-Net的架构，采用轴向深度可分离卷积，以减少参数量，同时保持特征提取能力。
• 在编码器和解码器路径中均采用7×7的轴向深度可分离卷积，以在不显著增加参数量的情况下扩大感受野。
• 引入双分支瓶颈结构，采用3×3轴向空洞深度可分离卷积，以增强特征表示并捕捉多尺度上下文信息。
• 优化网络架构，以最小化浮点运算量（FLOPs）和内存占用，同时保持高分割精度。
• 在标准医学图像分割基准数据集上，使用标准损失函数和数据增强技术，端到端训练并评估模型。
• 采用参数高效的架构设计理念，聚焦于深度可分离性和轴向卷积结构，以实现最小的参数量。

实验结果

研究问题

• RQ1 基于CNN的医学图像分割模型是否能在参数量少于100万的情况下实现最先进性能？
• RQ2 与标准卷积相比，使用轴向深度可分离卷积对分割精度和计算效率有何影响？
• RQ3 在轻量化架构中，瓶颈部分采用轴向空洞深度可分离卷积在多大程度上提升了特征表示能力？
• RQ4 在医学图像分割任务中，U-Lite与标准U-Net及现代基于Transformer的模型相比，在性能和效率方面表现如何？
• RQ5 是否可行设计一种高度参数高效的模型，同时保持在临床环境中的实际可部署性？

主要发现

• 在ACDC 2017数据集上，U-Lite的平均Dice分数达到92.1%，优于标准U-Net和多个最先进模型。
• 该模型仅包含87.8万个参数，比标准U-Net少35倍，同时保持或提升了分割精度。
• 与标准U-Net相比，U-Lite将浮点运算量（FLOPs）降低了90%以上，显著降低了计算成本。
• 采用轴向深度可分离卷积可在参数增加极少的情况下实现大感受野。
• 双分支瓶颈结构结合空洞深度可分离卷积，增强了特征表示，显著提升了在具有挑战性的医学图像分割任务中的性能。
• 该模型在不同医学影像模态和数据集上均表现出良好的泛化能力，表明其具备鲁棒性与实际可部署性。