[论文解读] OSS-Net: Memory Efficient High Resolution Semantic Segmentation of 3D Medical Data

Christoph Reich, Tim Prangemeier+2 more|arXiv (Cornell University)|Oct 20, 2021

Radiomics and Machine Learning in Medical Imaging被引用 7

一句话总结

OSS-Net 是一种内存高效、高分辨率的医学影像 3D 语义分割框架，通过引入局部观测和先验引导推理策略，对占用网络（O-Nets）进行扩展。其分割精度与 3D U-Net 基线相当，同时显著降低内存占用并实现更快的推理速度，使消费级 GPU 上的高分辨率 3D 分割成为可能。

ABSTRACT

Convolutional neural networks (CNNs) are the current state-of-the-art meta-algorithm for volumetric segmentation of medical data, for example, to localize COVID-19 infected tissue on computer tomography scans or the detection of tumour volumes in magnetic resonance imaging. A key limitation of 3D CNNs on voxelised data is that the memory consumption grows cubically with the training data resolution. Occupancy networks (O-Nets) are an alternative for which the data is represented continuously in a function space and 3D shapes are learned as a continuous decision boundary. While O-Nets are significantly more memory efficient than 3D CNNs, they are limited to simple shapes, are relatively slow at inference, and have not yet been adapted for 3D semantic segmentation of medical data. Here, we propose Occupancy Networks for Semantic Segmentation (OSS-Nets) to accurately and memory-efficiently segment 3D medical data. We build upon the original O-Net with modifications for increased expressiveness leading to improved segmentation performance comparable to 3D CNNs, as well as modifications for faster inference. We leverage local observations to represent complex shapes and prior encoder predictions to expedite inference. We showcase OSS-Net's performance on 3D brain tumour and liver segmentation against a function space baseline (O-Net), a performance baseline (3D residual U-Net), and an efficiency baseline (2D residual U-Net). OSS-Net yields segmentation results similar to the performance baseline and superior to the function space and efficiency baselines. In terms of memory efficiency, OSS-Net consumes comparable amounts of memory as the function space baseline, somewhat more memory than the efficiency baseline and significantly less than the performance baseline. As such, OSS-Net enables memory-efficient and accurate 3D semantic segmentation that can scale to high resolutions.

研究动机与目标

• 解决 3D 卷积神经网络（CNN）在高分辨率医学影像分割中内存呈立方增长的问题。
• 克服标准占用网络（O-Nets）在 3D 医学分割中表达能力有限且推理速度慢的局限。
• 将 O-Nets 的内存效率与 3D CNN 的性能相结合，用于 3D 语义分割。
• 通过降低内存消耗和推理时间，实现在消费级 GPU 上的高分辨率分割。
• 开发一种可扩展的、基于函数空间的方法，以捕捉脑瘤和肝脏病变等复杂解剖结构。

提出的方法

• OSS-Net 通过引入局部观测来增强对复杂 3D 形状的表征能力，对原始 O-Net 进行改进。
• 利用先验编码器预测作为初始状态，加速推理，使运行时间减少一倍。
• 模型采用可微分采样策略，每个前向传播过程评估多个 3D 位置的占用情况。
• 多尺度特征提取机制使网络能够捕捉体素数据中的细微解剖细节。
• 架构采用类似 3D U-Net 的编码器-解码器主干网络，专为函数空间表示而设计。
• 通过基于先验预测动态采样位置，优化推理过程，提升效率而不损失精度。

实验结果

研究问题

• RQ1 能否将占用网络扩展，使其在 3D 医学影像分割中的性能媲美 3D CNN？
• RQ2 如何提升 O-Nets 的表达能力，以建模脑瘤等复杂解剖结构？
• RQ3 能否在保持分割精度的前提下，显著加速 O-Nets 的推理速度？
• RQ4 在函数空间 3D 分割中，采样位置与内存消耗之间的最优权衡是什么？
• RQ5 OSS-Net 在训练数据有限的挑战性数据集（如 LiTS 肝脏分割数据集）上的表现如何？

主要发现

• 在 BraTS 2020 脑瘤数据集上，OSS-Net 的分割性能与 3D 残差 U-Net（性能基线）相当，Dice 分数达到 0.878。
• 在 LiTS 肝脏分割数据集上，OSS-Net 超过 O-Net 基线和 2D U-Net 效率基线，Dice 分数达到 0.852。
• 与原始 O-Net 方法相比，推理速度提升两倍，且在采样位置超过 2^12 个后，运行时间保持恒定。
• 推理期间的内存使用量与 O-Net 相当（1.2GB），且比 3D U-Net 基线（6.5GB）低五倍以上，支持在消费级 GPU 上部署。
• 原始方法与改进后推理方法之间的预测差异仅为 5.9658×10⁻⁴%，证实性能下降可忽略不计。
• OSS-Net 训练期间使用 2.6GB 内存，显著低于 3D U-Net 基线（6.5GB），同时保持高精度。