[논문 리뷰] End-to-end training of a two-stage neural network for defect detection

Jakob Božič, Domen Tabernik외 1명|arXiv (Cornell University)|2020. 07. 15.

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%로 향상시켰다.
• 절단 실험 결과, 동적 손실 균형 조정, 기울기 흐름 조정, 빈도 사용 샘플링, 거리변환 가중치 등 각각의 제안된 구성 요소가 성능 향상에 개별적으로 기여하며, 모든 구성 요소를 함께 사용했을 때 최고의 성능가를 달성함을 확인하였다.
• 이 방법은 군더미 레이블에 대해 강건하며, 정밀한 픽셀 수준의 마스크가 아닌 영역 기반 레이블을 사용해도 효과적인 학습이 가능함을 보여주었다. 이는 정밀한 레이블링이 비용이 많이 들기 때문에 중요하다.
• 절단 실험 결과, 구성 요소 중 하나를 제거할 경우 성능 저하가 명백하게 관찰되었으며, 이는 제안된 모든 확장 기능이 상호 보완적이며 필수적임을 확인한다.