[论文解读] PyTorch Connectomics: A Scalable and Flexible Segmentation Framework for EM Connectomics

Zudi Lin, Donglai Wei+2 more|arXiv (Cornell University)|Dec 10, 2021

Cell Image Analysis Techniques被引用 27

一句话总结

PyTorch Connectomics (PyTC) 是一个用于体积 EM 图像的语义与实例分割的开源框架，能够实现可扩展的数据加载、灵活的模型架构、多目标损失和半监督/自监督学习，且在 CREMI 突触沟裂结果方面达到了最先进水平。

ABSTRACT

We present PyTorch Connectomics (PyTC), an open-source deep-learning framework for the semantic and instance segmentation of volumetric microscopy images, built upon PyTorch. We demonstrate the effectiveness of PyTC in the field of connectomics, which aims to segment and reconstruct neurons, synapses, and other organelles like mitochondria at nanometer resolution for understanding neuronal communication, metabolism, and development in animal brains. PyTC is a scalable and flexible toolbox that tackles datasets at different scales and supports multi-task and semi-supervised learning to better exploit expensive expert annotations and the vast amount of unlabeled data during training. Those functionalities can be easily realized in PyTC by changing the configuration options without coding and adapted to other 2D and 3D segmentation tasks for different tissues and imaging modalities. Quantitatively, our framework achieves the best performance in the CREMI challenge for synaptic cleft segmentation (outperforms existing best result by relatively 6.1$\%$) and competitive performance on mitochondria and neuronal nuclei segmentation. Code and tutorials are publicly available at https://connectomics.readthedocs.io.

研究动机与目标

• 为 EM connectomics 数据集（不同大小与模态）提供一个可扩展、灵活的分割工具箱的需求。
• 提供一个模块化框架，实现在神经元、突触、线粒体和细胞核之间的语义与实例分割。
• 支持多任务学习、半监督/自监督学习，以及可配置的架构，以利用有限标签和大量未标注数据。
• 通过配置文件实现易于定制，以处理各向异性的 2D/3D 数据及多样的成像模态。

提出的方法

• 模块化四组件系统：数据模块用于可扩展的数据加载（Volume Dataset 与 Tile Dataset），模型模块用于支持多目标的 2D/3D 或混合架构，Engine 模块用于分布式/非分布式训练和多 GPU 推理，配置模块（YACS）用于可重复性。
• 混合表示学习，允许从单一共享网络中以多种损失学习多个目标（例如前景掩码、实例轮廓、带符号距离变换）。
• 主动学习与半监督学习工作流（两流主动查询；带伪标签的自训练）以充分利用未标注数据。
• 灵活的网络架构通过阶段性先 2D 再 3D 卷积实现对各向异性数据的处理，支持 U-Net、DeepLab 风格的骨干网以及混合 2D/3D 卷积。
• 数据增强包括非时空的（对比度、亮度、伽玛、缺失部分）和时空的（重采样、错位），以及测试时增强（多达 16 种变体）。
• 分布式训练可使用 DataParallel 或 DistributedDataParallel，非重叠数据分区以平衡前景密度，推理工作流将大体积切分以实现并行处理。

实验结果

研究问题

• RQ1 单一框架如何在可配置的多目标损失下同时支持 EM connectomics 中的语义与实例分割任务？
• RQ2 可扩展的数据加载与分布式训练是否能够实现对 PB 级 EM 数据集的实际分割？
• RQ3 在标注数据稀缺时，混合表示学习与半监督/自监督学习是否提升分割性能？
• RQ4 对 nm 分辨率的 nanometer EM 数据而言，各向异性 2D/3D 架构和混合卷积的有效性如何？
• RQ5 主动学习与自训练在具有挑战性的 connectomics 分割任务中的增益如何？

主要发现

• PyTC 实现了 CREMI 突触沟裂分割的最先进水平，半监督学习比先前最佳结果提升了总体 CREMI 得分 6.1%。
• 在突触沟裂数据上，SSL 方法在 CREMI 提交中排名第一，并在 ADGT、ADF 和 CREMI 指标上优于现有方法。
• 在线粒体语义分割方面，基于 PyTC 的模型超过了先前方法，甚至在简单后处理下接近或超越人工标注性能。
• 在线粒体实例分割（MitoEM）中，U3D-BC/U3D-BCD 模型取得具有竞争力的 AP-75 分数，在挑战参与者中排名第三，凸显了在大规模实例分割中的有效性。
• 在神经元核的分割方面，U3D-BC 与 U3D-BCD 变体优于若干基线方法（Cellpose、Stardist），距离映射变体在整体 AP 上带来显著提升（约 22% 的提升）。
• 跨数据集，PyTC 展现了可扩展的数据处理能力（Volume 与 Tile 数据集）、多 GPU 训练/推理，以及对低标签场景有益的灵活多目标学习。