用可微分选区取代滑动窗口,让3D医学图像分割更快更省显存
No More Sliding Window: Efficient 3D Medical Image Segmentation with Differentiable Top-k Patch Sampling
- 通过可微Top-k采样,只选关键切片块,跳过传统滑动窗口
- 在保持精度的同时,计算量降为原来的9%,推理速度提升9~11倍
- 适用于任意3D分割模型,特别适合资源受限场景使用
3D模型在CT/MRI分割中因能捕捉跨层关系而优于2D模型,但深度维度显著增加内存消耗。尽管基于切片块的训练缓解了内存压力,但滑动窗口(SW)方法严重拖慢推理速度。本文提出无滑动窗口(NMSW)框架,通过可微分Top-k模块仅选取最相关切片块,消除对SW的需求。当切片级预测不足时,系统智能利用粗粒度全局预测进行结果修正。在三个任务和三种分割骨干网络上评估,NMSW在保持竞争力精度的同时,将计算复杂度从88.0降至8.00 TMACs(降低91%),在H100 GPU上推理时间由99.0秒缩短至8.3秒(提速9.1倍),在Xeon Gold CPU上由2110秒降至189秒(提速11.1倍)。NMSW具有模型无关性,与现有高效分割骨干结合后效率进一步提升。代码已开源:https://github.com/Youngseok0001/open_nmsw。
原文摘要 · Abstract (English)
3D models surpass 2D models in CT/MRI segmentation by effectively capturing inter-slice relationships. However, the added depth dimension substantially increases memory consumption. While patch-based training alleviates memory constraints, it significantly slows down the inference speed due to the sliding window (SW) approach. We propose No-More-Sliding-Window (NMSW), a novel end-to-end trainable framework that enhances the efficiency of generic 3D segmentation backbone during an inference step by eliminating the need for SW. NMSW employs a differentiable Top-k module to selectively sample only the most relevant patches, thereby minimizing redundant computations. When patch-level predictions are insufficient, the framework intelligently leverages coarse global predictions to refine results. Evaluated across 3 tasks using 3 segmentation backbones, NMSW achieves competitive accuracy compared to SW inference while significantly reducing computational complexity by 91% (88.0 to 8.00 TMACs). Moreover, it delivers a 9.1x faster inference on the H100 GPU (99.0 to 8.3 sec) and a 11.1x faster inference on the Xeon Gold CPU (2110 to 189 sec). NMSW is model-agnostic, further boosting efficiency when integrated with any existing efficient segmentation backbones. The code is avaialble: https://github.com/Youngseok0001/open_nmsw.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。