[論文レビュー] 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

ひとこと要約

本稿では、表面欠陥検出における2段階ニューラルネットワークのエンドツーエンド学習スキームを提案し、勾配フロー調整と動的損失バランス化により、セグメンテーションと分類の同時最適化を可能にした。本手法は、頻度ベースサンプリングと距離変換重み付き損失を用いることで、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.

研究の動機と目的

• セグメンテーションと分類ヘッドの順次学習を必要とする、欠陥検出ネットワークにおける遅く面倒な2段階学習プロセスを解決すること。
• パフォーマンスを損なわず、2段階アーキテクチャでエンドツーエンド学習を可能にし、訓練時間を短縮するとともに収束性を向上させること。
• 正確なピクセル単位のアノテーションへの依存度を低減するため、不確実性モデリングを組み込んだ損失関数の拡張により、領域ベースの粗いアノテーションを扱えるようにすること。
• 負例（非欠陥）サンプルの頻度ベースサンプリングを導入することで、クラス不均衡なデータセットにおけるモデル一般化性能を向上させること。
• 正例ピクセルの重み付けに距離変換を適用することで、欠陥領域の特徴学習を強化し、高確率の欠陥領域に注目を集中させること。

提案手法

• 1回のフォワードおよびバックワードパス内でセグメンテーションと分類ヘッドを同時に最適化することでエンドツーエンド学習を導入し、順次学習の必要性を排除すること。
• バックプロパゲーション中にセグメンテーションと分類の損失成分を動的にバランス化することで、分類勾配による不安定性を防ぎ、学習を安定化させること。
• 共同最適化中の学習済み特徴を損なわないように、分類ヘッドからセグメンテーションヘッドへの勾配フローを調整すること。
• 領域ベースのアノテーションにおける不確実性を考慮するように損失関数を拡張し、正確さに劣るがアノテーションが容易なバウンディングボックスや領域を活用できるようにすること。
• 負例（非欠陥）サンプルの使用頻度が低いものを優先する頻度ベースサンプリング戦略を適用し、訓練中のクラス不均衡を是正すること。
• 正例ピクセルの重み付けにセグメンテーションマスクの距離変換をクラスに依存しない重みとして用い、欠陥境界付近のピクセルに高い重要性を与えることで、局所化精度を向上させること。

実験結果

リサーチクエスチョン

• RQ1 エンドツーエンド学習による2段階欠陥検出ネットワークは、別々の最適化フェーズを経る従来の2段階学習と比較して優れた性能を達成できるか？
• RQ2 共同最適化中のセグメンテーションヘッドと分類ヘッド間の勾配フローをどのように調整すれば、安定した特徴学習を維持できるか？
• RQ3 検出精度を損なわず、領域ベースの粗いアノテーションをどれだけ効果的に使用できるか？
• RQ4 負例サンプルの頻度ベースサンプリングは、クラス不均衡な欠陥検出データセットにおけるモデル性能を向上させるか？
• RQ5 距離変換に基づく正例ピクセルの重み付けは、正確なピクセル単位のアノテーションが不要な状況でも検出性能を向上させられるか？

主な発見

• 提案されたエンドツーエンド学習スキームは、DAGMデータセットで100%の検出率を達成し、完全にこのデータセットを解消し、すべての先行手法を上回った。
• KolektorSDDデータセットでは、100%の検出率を達成し、同様にデータセットを完全に解消したが、従来の2段階アプローチと比較して訓練イテレーションを半分に削減した。
• 頻度ベースサンプリングと距離変換重み付けの組み合わせにより、Severstal Steel欠陥データセットのAPが98.24%から98.70%に向上した。
• アブレーションスタディの結果、動的損失バランス化、勾配フロー調整、頻度ベースサンプリング、距離変換重み付けという各提案コンポーネントが個別に性能向上に寄与していることが確認され、すべてを組み合わせた場合に最高の結果が得られた。
• 本手法は粗いアノテーションに対してロバストであることが示され、正確なピクセル単位のマスクではなく領域ベースのラベルでも効果的な学習が可能になった。
• アブレーションスタディの結果、1つのコンポーネントを除去すると性能が顕著に低下することが確認され、提案された拡張の必要性と相乗効果の有効性が裏付けられた。