在3D多类分割中扩散模型训练策略与评估对齐的重要性
图像与视频处理
2023-08-21 v3 计算机视觉与模式识别
摘要
近来，去噪扩散概率模型（DDPM）已被应用于图像分割，通过以图像为条件生成分割掩码，但其应用主要局限于2D网络，未利用3D公式的潜在优势。在本工作中，我们研究了基于DDPM的分割模型在两个大型多类数据集（前列腺MR和腹部CT）上的3D多类分割。我们观察到，训练与测试方法之间的差异导致现有DDPM方法性能不佳。为缓解这种不一致性，我们提出了一种循环方法，该方法基于模型在前一时间步的预测而非使用真值来生成 corrupted masks（ corrupted 掩码）。所提方法相较于现有DDPM取得了统计显著的性能提升，且独立于其他多种减少训练-测试差异的技术，包括执行掩码预测、使用Dice损失以及在训练期间减少扩散时间步数。扩散模型的性能也与基于非扩散的U-net具有竞争力且视觉上相似，在相同计算预算内。基于JAX的扩散框架已发布于 https://github.com/mathpluscode/ImgX-DiffSeg。
引用
@article{arxiv.2303.06040,
title = {Importance of Aligning Training Strategy with Evaluation for Diffusion Models in 3D Multiclass Segmentation},
author = {Yunguan Fu and Yiwen Li and Shaheer U. Saeed and Matthew J. Clarkson and Yipeng Hu},
journal= {arXiv preprint arXiv:2303.06040},
year = {2023}
}
备注
Accepted at Deep Generative Models workshop at MICCAI 2023