OctreeNCA:在消费级硬件上单次处理 184 MP 分割
计算机视觉与模式识别
2025-08-12 v1
摘要
医学应用要求分割大尺寸输入,例如前列腺 MRI、病理切片或手术视频。理想情况下,这些输入应一次性推理,以便为模型提供适当的空间或时间上下文。在分割大尺寸输入时,GPU 的显存消耗成为瓶颈。UNet 或 Vision Transformer 等架构的显存消耗扩展性很差,导致采用基于图像块或逐帧的方法,这损害了全局一致性和推理速度。轻量级的神经细胞自动机(NCA)是一种受生物启发的模型,其构造本身具有尺寸不变性。然而,由于其仅限局部的通信规则,它缺乏全局知识。我们通过使用八叉树数据结构推广邻域定义,提出了 OctreeNCA。我们推广的邻域定义能够高效地遍历全局知识。由于深度学习框架主要针对大型多层网络开发,其实现并未充分利用 NCA 的优势。我们在 CUDA 中实现了一个 NCA 推理函数,进一步降低了显存需求并提高了推理速度。我们的 OctreeNCA 能够快速分割高分辨率图像和视频,同时在评估期间比 UNet 少占用 90% 的显存。这使我们能够一次性分割 184 兆像素的病理切片或 1 分钟的手术视频。
引用
@article{arxiv.2508.06993,
title = {OctreeNCA: Single-Pass 184 MP Segmentation on Consumer Hardware},
author = {Nick Lemke and John Kalkhof and Niklas Babendererde and Anirban Mukhopadhyay},
journal= {arXiv preprint arXiv:2508.06993},
year = {2025}
}