跳转至

Revisiting nnU-Net for Iterative Pseudo Labeling and Efficient Sliding Window Inference

论文地址

https://openreview.net/forum?id=FNMbe2vLvev

代码

https://github.com/Ziyan-Huang/FLARE22

摘要

Cite

nnU-Net作为近年来许多医学图像分割挑战的良好基线。它适用于完全监督的分割任务。然而，它的推理效率较低，并且不能有效地充分利用未标记的数据， 这两者在实际临床场景中都是至关重要的。为此，我们重新审视nnU-Net，并在此框架中找到效率和准确性之间的权衡。基于默认的nnU-Net设置， 我们设计了一个由两种策略组成的协同训练框架，分别生成高质量的伪标签和进行有效的推理。具体来说，我们首先设计了一个资源密集型的nnU-Net 来为未标记的数据迭代生成高质量的伪标签。然后我们使用标记数据和选择的未标记数据训练另一个轻量级3D nnU-Net，后者使用高质量伪标签以实现有效分割。 我们对FLARE22挑战进行了实验。我们的资源密集型nnU-Net在13个腹部器官分割任务上实现了0.9064的平均DSC，并在验证排行榜上名列第一。 我们的轻量级nnUNet在验证排行榜上显示平均DSC为0.8773，但它在准确性和效率之间取得了更好的权衡。在测试集上，它显示平均DSC为 0.8864， 平均NSD 为 0.9465，平均推理时间为 14.59s，赢得了FLARE22 挑战赛的冠军。我们的代码可在 https://github.com/Ziyan-Huang/FLARE22 上公开获得。

对nnU-Net的拓展，主要尝试解决nnU-Net的两个痛点，一个是无法使用未标注的数据，另一个是推理效率低。FLARE22比赛是分割腹部的13个器官，给了50个标注的数据， 还有2000个没有标注的数据，这种数据分布是符合医学图像处理的数据规律的，因此这种情况下如何使用没有标注的数据就变得十分重要。

方法

本文的贡献包括

  1. 设计了一个基于 nnU-Net 的伪标签框架，可以生成高质量的伪标签并同时高效地进行推理。
  2. 提出了一种图像级伪标签选择方法，该方法基于不同重新训练迭代期间伪标签的稳定性。使用选择的伪标签训练的模型表现更好。
  3. 通过考虑腹部器官体积的先验知识提出了一种有效的滑动窗口推理策略。这种策略可以大大减少推理窗口的数量。
  4. 优化了nnU-Net中代码的耗时部分，例如crop和resample。


流程图还是比较简单明了的，整个pipeline流程是

  1. 用标注的数据训练一个大的UNet（五折交叉验证）
  2. 用大的UNet来推理未标注的数据，得到伪标签
  3. 用标注的数据和伪标签继续训练大的UNet，然后再分割，继续得到伪标签
  4. 通过判定分割的稳定性，以及步骤3和步骤2的结果对比，优化步骤3中获得的伪标签，这个步骤可能迭代多次
  5. 使用标注的数据和步骤4中得到的效果较好的伪标签，训练一个小的UNet，来加速推理性能


优化的滑动窗口推理，没有看的特别懂，大概意思就是由于有先验知识，比如胃肯定在人体的中间附近，那我就从中间竖着一条推理试试，有前景我就开始分周围的， 没有的话这一个区域就都不推理了，出发点非常好理解，但是太依赖于先验了，我要是有这么强的先验我就直接crop了。

实验结果



看实验结果是非常香的，加上伪标签后不管是大模型还是小模型效果都有提升，在一些种类的数据集上，小模型+伪标签甚至超过了不加伪标签的大模型的效果。 另一个现象是伪标签对小模型的提升较大

总结

Cite

在本文中，我们设计了一个基于nnU-Net的框架，以使用未标记的数据进行训练并有效地进行推理。我们相信我们提出的框架可以作为半监督学习和医学图像分割有效推理的良好基线。

主要研究的是在有限标记数据下，并且要求使用小模型，如何提升性能

评论