[论文解读] End-to-end training of a two-stage neural network for defect detection

Jakob Božič, Domen Tabernik+1 more|arXiv (Cornell University)|Jul 15, 2020

Industrial Vision Systems and Defect Detection被引用 1

一句话总结

本文提出了一种用于表面缺陷检测中两阶段神经网络的端到端训练方案，通过梯度流调节与动态损失平衡，实现分割与分类的联合优化。该方法通过频率使用采样和基于距离变换的正样本损失加权等创新手段，实现了最先进性能，在 DAGM 与 KolektorSDD 数据集上达到 100% 检测率。

ABSTRACT

Segmentation-based, two-stage neural network has shown excellent results in the surface defect detection, enabling the network to learn from a relatively small number of samples. In this work, we introduce end-to-end training of the two-stage network together with several extensions to the training process, which reduce the amount of training time and improve the results on the surface defect detection tasks. To enable end-to-end training we carefully balance the contributions of both the segmentation and the classification loss throughout the learning. We adjust the gradient flow from the classification into the segmentation network in order to prevent the unstable features from corrupting the learning. As an additional extension to the learning, we propose frequency-of-use sampling scheme of negative samples to address the issue of over- and under-sampling of images during the training, while we employ the distance transform algorithm on the region-based segmentation masks as weights for positive pixels, giving greater importance to areas with higher probability of presence of defect without requiring a detailed annotation. We demonstrate the performance of the end-to-end training scheme and the proposed extensions on three defect detection datasets - DAGM, KolektorSDD and Severstal Steel defect dataset - where we show state-of-the-art results. On the DAGM and the KolektorSDD we demonstrate 100\% detection rate, therefore completely solving the datasets. Additional ablation study performed on all three datasets quantitatively demonstrates the contribution to the overall result improvements for each of the proposed extensions.

研究动机与目标

• 解决缺陷检测网络中两阶段训练过程缓慢且繁琐的问题，该过程需要对分割头和分类头进行顺序训练。
• 在不损害性能的前提下，实现在两阶段架构中的端到端学习，减少训练时间并改善收敛性。
• 通过扩展损失函数以处理具有不确定性的区域级、粗粒度标注，降低对精确像素级标注的依赖。
• 通过引入频率使用采样策略，提升模型在类别不平衡数据集上的泛化能力，优先处理使用频率较低的负样本（非缺陷样本）。
• 通过将距离变换应用于分割掩码作为正样本的类无关权重，增强对缺陷区域的特征学习，尤其提升对高概率缺陷区域的关注。

提出的方法

• 通过在一次前向与反向传播中联合优化分割头与分类头，实现端到端训练，消除对顺序训练的需求。
• 在反向传播过程中动态平衡分割与分类损失分量，以稳定训练过程，防止分类梯度导致的不稳定性。
• 调整从分类头到分割头的梯度流，避免在联合优化过程中破坏已学习的特征。
• 扩展损失函数以考虑区域级标注中的不确定性，从而支持使用更易标注但精度较低的边界框或区域标注。
• 应用频率使用采样策略，优先选择使用频率较低的负样本（非缺陷样本），以缓解训练过程中的类别不平衡问题。
• 将距离变换应用于分割掩码，作为正样本的类无关权重，使靠近缺陷边界的像素获得更高重要性，从而提升定位精度。

实验结果

研究问题

• RQ1 与传统两阶段训练中分阶段优化的方法相比，两阶段缺陷检测网络的端到端训练是否能实现更优性能？
• RQ2 在联合优化过程中，如何调节分割头与分类头之间的梯度流，以维持稳定的特征学习？
• RQ3 在不牺牲检测准确率的前提下，能在多大程度上有效使用更粗粒度的区域级标注进行训练？
• RQ4 负样本的频率使用采样是否能提升模型在类别不平衡缺陷检测数据集上的性能？
• RQ5 基于距离变换的正样本加权方法是否能在无需精确像素级标注的情况下提升检测性能？

主要发现

• 所提出的端到端训练方案在 DAGM 数据集上实现了 100% 检测率，完全解决了该数据集，优于所有先前方法。
• 在 KolektorSDD 数据集上，该方法也实现了 100% 检测率，完全解决了该数据集，并且相比先前的两阶段方法将训练迭代次数减少了一半。
• 频率使用采样与距离变换加权的结合使 Severstal Steel 缺陷数据集上的 AP 从 98.24% 提升至 98.70%。
• 消融实验表明，所提出的各项组件——动态损失平衡、梯度流调节、频率使用采样与距离变换加权——各自均对性能提升有贡献，当所有组件同时使用时达到最佳效果。
• 该方法对粗粒度标注具有鲁棒性，能够有效使用区域级标签进行训练，而无需昂贵的像素级掩码。
• 消融实验表明，移除任意单一组件均会导致性能明显下降，证实了所有所提扩展方法的必要性及其相互间的协同增益作用。