[论文解读] Recurrent U-Net for Resource-Constrained Segmentation

Wei Wang, Kaicheng Yu+3 more|arXiv (Cornell University)|Jun 11, 2019

Advanced Neural Network Applications参考文献 37被引用 12

一句话总结

本文提出一种循环U-Net架构，通过在多个层级上迭代优化分割掩码和网络内部状态，提升分割性能，在手部、视网膜血管和道路分割任务上实现最先进结果，仅使用0.3百万参数，优于DeepLabV3+和RefineNet等更大模型在资源受限数据集上的表现。

ABSTRACT

State-of-the-art segmentation methods rely on very deep networks that are not always easy to train without very large training datasets and tend to be relatively slow to run on standard GPUs. In this paper, we introduce a novel recurrent U-Net architecture that preserves the compactness of the original U-Net, while substantially increasing its performance to the point where it outperforms the state of the art on several benchmarks. We will demonstrate its effectiveness for several tasks, including hand segmentation, retina vessel segmentation, and road segmentation. We also introduce a large-scale dataset for hand segmentation.

研究动机与目标

• 解决深度大参数分割模型在AR头戴设备等资源受限环境中的局限性。
• 在不增加模型大小或依赖大规模预训练数据集的前提下，提升紧凑U-Net架构的性能。
• 实现在小规模专用数据集上实现高精度的实时推理，避免大模型过拟合。
• 引入一个新的大规模手部分割数据集，以支持低数据场景下的基准测试。
• 证明在多个U-Net层级上进行循环优化优于仅更新单个隐藏层的简单循环架构和标准U-Net。

提出的方法

• 引入一种跨多个编码器和解码器层的循环单元（DRU），实现对内部表征的迭代优化。
• 通过在循环迭代中共享参数，保持低内存占用和计算成本。
• 采用一种反馈循环，同时更新分割掩码和网络内部隐藏状态。
• 在多层上应用循环机制，使高层特征可被迭代优化，区别于单层RNN。
• 从零开始在小规模数据集上训练模型，避免依赖ImageNet预训练。
• 采用深度为4的残差U-Net主干网络构建DRU，实现高效的特征传播与优化。

实验结果

研究问题

• RQ1 在不增加模型大小的前提下，循环U-Net架构是否能提升小数据集上的分割精度？
• RQ2 在多个U-Net层级上进行循环优化是否优于仅更新单个隐藏层的简单循环基线模型？
• RQ3 仅含0.3百万参数的紧凑模型是否能在专用任务上超越DeepLabV3+和RefineNet等大型预训练模型？
• RQ4 该方法在训练样本有限的数据集（如手部和视网膜血管分割）上的表现如何？
• RQ5 循环U-Net是否能在消费级硬件上实现高精度的实时推理？

主要发现

• 在新构建的Keyboard Hand Dataset（12.5K张图像）上，所提出的循环U-Net在Titan X GPU上实现55帧每秒的推理速度，证明其具备实时性能。
• 在DRIVE视网膜血管分割数据集上，模型mIoU达到0.821，优于此前最先进方法（F1分数0.822 vs. 0.92，本方法更高）。
• 在Cityscapes验证集上，DRU-VGG16变体达到mIoU 0.761，性能与使用ResNet101主干的DeepLabV3+相当。
• 在道路分割任务上，所提方法mIoU达0.560，较U-Net-G高出8个百分点，较DRU-VGG16高出5个百分点。
• 在速度-精度权衡上优于ICNet，推理速度更快且精度显著更高。
• 在五个基准测试中，包括新构建的12.5K图像数据集，该方法在手部分割任务上表现优于基于RefineNet的最先进方法。