提出轻量级Unet,参数减少30倍仍保持分割精度
Lean Unet: A Compact Model for Image Segmentation
- 采用扁平结构,各层通道数不变,避免传统下采样时通道翻倍
- 在MRI和CT数据集上性能接近甚至优于原版Unet和剪枝模型
- 无需训练剪枝,适合医疗影像等资源受限场景
Unet及其变体在语义图像分割中广泛应用,尤其在计算机辅助放射学领域。现有Unet架构通过逐层降低空间分辨率并增加通道数来保持信息内容,但导致内存占用大,限制了训练批次大小并增加了推理延迟。通道剪枝可压缩Unet而保持精度,但需长时间优化且泛化能力差。通过对剪枝过程的分析,我们发现最终结构比剪枝策略更关键。基于此,提出一种紧凑的LUnet架构,其具有扁平层次结构,不随分辨率减半而翻倍通道数。在公开MRI数据集及两个内部CT数据集上评估,结果显示先进剪枝方法(STAMP)主要剪除通道数最多的层;相比之下,随机移除这些层中任一通道或最大通道层即能达到相似甚至更好性能。所提LUnet采用固定架构,参数量超30倍减少,性能仍与标准Unet及数据自适应剪枝网络相当。固定通道数的LUnet在相同参数量下表现优于标准Unet。跳跃连接使瓶颈通道大幅减少,不同于传统编码器-解码器架构需增加瓶颈通道以传递信息。
原文摘要 · Abstract (English)
Unet and its variations have been standard in semantic image segmentation, especially for computer assisted radiology. Current Unet architectures iteratively downsample spatial resolution while increasing channel dimensions to preserve information content. Such a structure demands a large memory footprint, limiting training batch sizes and increasing inference latency. Channel pruning compresses Unet architecture without accuracy loss, but requires lengthy optimization and may not generalize across tasks and datasets. By investigating Unet pruning, we hypothesize that the final structure is the crucial factor, not the channel selection strategy of pruning. Based on our observations, we propose a lean Unet architecture (LUnet) with a compact, flat hierarchy where channels are not doubled as resolution is halved. We evaluate on a public MRI dataset allowing comparable reporting, as well as on two internal CT datasets. We show that a state-of-the-art pruning solution (STAMP) mainly prunes from the layers with the highest number of channels. Comparatively, simply eliminating a random channel at the pruning-identified layer or at the largest layer achieves similar or better performance. Our proposed LUnet with fixed architectures and over 30 times fewer parameters achieves performance comparable to both conventional Unet counterparts and data-adaptively pruned networks. The proposed lean Unet with constant channel count across layers requires far fewer parameters while achieving performance superior to standard Unet for the same total number of parameters. Skip connections allow Unet bottleneck channels to be largely reduced, unlike standard encoder-decoder architectures requiring increased bottleneck channels for information propagation.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。