用哈特利变换设计极小模型,实现高效高分辨率3D医学图像分割。
HNOSeg-XS: Extremely Small Hartley Neural Operator for Efficient and Resolution-Robust 3D Image Segmentation
- 用哈特利变换替代傅里叶变换,构建频率域的微分方程学习模型。
- 参数少于3.47万,在多个数据集上推理时间低于0.24秒,内存占用小于1.8 GiB。
- 无需重新训练即可处理更高分辨率图像,适合资源受限场景下的医学影像分析。
在医学图像分割中,卷积神经网络(CNN)和变换器(transformer)是主流方法。CNN依赖连续卷积和池化捕获长距离空间相关性,但计算开销大,3D模型层数受限;transformer虽能捕捉长程依赖,但其复杂度随输入大小呈二次增长,计算成本高。因此,两类模型常需降低输入尺寸以增加层数与滤波器数量,但导致训练与推理分辨率不一致,影响高分辨率表现。为此,本文提出新型分辨率鲁棒的HNOSeg-XS架构。通过傅里叶神经算子建模可学习的偏微分方程,并将傅里叶变换替换为哈特利变换,实现频率域重定义。该模型具备零样本超分辨率能力,具有极低参数量、快速推理与内存高效等优势。在BraTS'23、KiTS'23与MVSeg'23数据集上,使用Tesla V100 GPU测试,模型参数少于34.7k,整体推理时间低于0.24秒,内存占用小于1.8 GiB,显著优于对比的CNN与transformer模型。
原文摘要 · Abstract (English)
In medical image segmentation, convolutional neural networks (CNNs) and transformers are dominant. For CNNs, given the local receptive fields of convolutional layers, long-range spatial correlations are captured through consecutive convolutions and pooling. However, as the computational cost and memory footprint can be prohibitively large, 3D models can only afford fewer layers than 2D models with reduced receptive fields and abstract levels. For transformers, although long-range correlations can be captured by multi-head attention, its quadratic complexity with respect to input size is computationally demanding. Therefore, either model may require input size reduction to allow more filters and layers for better segmentation. Nevertheless, given their discrete nature, models trained with patch-wise training or image downsampling may produce suboptimal results when applied on higher resolutions. To address this issue, here we propose the resolution-robust HNOSeg-XS architecture. We model image segmentation by learnable partial differential equations through the Fourier neural operator which has the zero-shot super-resolution property. By replacing the Fourier transform by the Hartley transform and reformulating the problem in the frequency domain, we created the HNOSeg-XS model, which is resolution robust, fast, memory efficient, and extremely parameter efficient. When tested on the BraTS'23, KiTS'23, and MVSeg'23 datasets with a Tesla V100 GPU, HNOSeg-XS showed its superior resolution robustness with fewer than 34.7k model parameters. It also achieved the overall best inference time (< 0.24 s) and memory efficiency (< 1.8 GiB) compared to the tested CNN and transformer models.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。