[논문 리뷰] End-to -End Semi-Supervised Object Detection with Soft Teacher

 

Introduce

본 논문은 ICCV2021에 accepted된 논문으로 semi-supervised learning을 이용한 object detection 모델이다.

현재 paperswithcode에서 COCO 데이터셋 기준 SOTA인 모델이다.

 

https://arxiv.org/abs/2106.09018

End-to-End Semi-Supervised Object Detection with Soft Teacher

 

 

Semi-Supervised Learning

일반적으로 데이터(의료, 이미지, 등)를 annotation하는데는 많은 시간과 비용(인적, 물적)이 발생한다.

이러한 문제를 완화시키기 위한 방법으로 사용되는 것이 semi supervision이다.

 

semi-supervised learning은 labeled 데이터와 unlabeled 데이터 모두 사용하는 방법이다.

이러한 semi-supervised에는 크게 두가지 방법이 있는데, Consistency methods 와 Pseudo-label method이다.

Consistency 방법의 경우, unlabeled 이미지를 활용하여 동일한 이미지에 다른 약간의 변화를 줘서 유사한 예측을 생성하는 방법으로 예를들어 하얀 고양이와 빨간 고양이 이미지가 있을때 둘 다 고양이 이미지라는 것을 알 수 있다.

Pseudo-label 방법의 경우, 초기에 훈련된 모델에 의한 pseudo labels로 unlabeled 이미지에 annotation하고, 이러한 pseudo labeled 이미지를 이용해서 detector를 refine한다.

 

해당 논문의 경우 pseudo-label 방법에 속한다.

 

 

Contribution

본 논문은 크게 4가지의 contribution이 있다.

 

1. End-to-End Pseudo Labeling Framework

• 학습을 진행하면서 점진적으로 pseudo label의 퀄리티를 향상시키고, 이를 통해 object detection의 성능을 향상시킨다.

 

2. Soft Teacher

• teacher 모델이 생성하는 pseudo label의 class의 reliability를 평가하는데 사용된다.

 

3. box jittering

• pseudo label의 box regression의 reliability를 평가하는데 사용된다.

 

4. SOTA

• paperswithcode에서 COCO 데이터셋 기준 상위권에 현재 랭크되어있다.

 

 

Overall

 

 

다음은 논문의 프레임워크 개요이다.

간단하게 설명하자면, 먼저 학습이 시작되면서 soft teacher 모델과 student 모델 둘 다 랜덤하게 초기화된다.

그리고 teacher 모델의 경우 student 모델의 EMA를 따라서 지속적으로 업데이트된다. 이는 mean teacher 라는 논문에서 가져온 방법이다.

이미지의 왼쪽을 보면 unlabeled 데이터를 사용할때 weak augmentation과 strong augmentation 사용하는데 이는 FixMatch라는 논문에서 가져온 방법으로 hight 퀄리티의 pseudo label를 만들기 위해 사용되었다.

unlabeled 이미지가 입력되었을때, teacher가 detection 하여 후보 box를 여러개 생성하고, 이를 NMS를 통해 중복을 제거한다. 그리고 그래도 남은 수많은 box 중 score가 threshold 보다 낮은 box는 추가 제거하여 classification loss에 사용할 pseudo label을 생성한다.

그리고 후보 박스를 box regression variance filter를 이용해서 regression loss를 위한 pseudo box를 생성한다. 이 부분이 앞에서 언급한 box jittering 방법이다.

마지막으로 전체 loss는 labeled 데이터를 이용해서 생성한 supervised loss와 unlabeled 데이터를 이용해서 생성한 loss를 모두 더해서 사용한다.

 

 

EMA(Exponential Moving Average)

 

다음은 앞에서 언급한 EMA로 student 모델을 이용하여 teacher 모델을 업데이트할때 사용한 방법이다.

이 방법은 mean teacher 라는 논문에서 가져온것으로 student 모델의 체크포인트의 최근 가중치에 높은 가중치를 적용하여 평균을 계산하여 teacher를 업데이트한다. 즉, 최근값에 높은 가중치를 적용하여 평균을 구하는 방법이다.

 

 

Soft Teacher

Motivation

일반적으로 detector의 성능은 pseudo label의 퀄리티에 따라 달라진다. 그런데 저자가 실험해본 결과 threshold(foreground score)를 0.9로 설정했을때, precision은  89%로 높았지만, recall은 33%로 낮았다. 즉, 일부 box가 잘못 지정된 경우가 발생할 수 있고 이는 학습을 방해해서 성능에 해(harm)가 될 수 있다.

 

참고로,

precision: 모델이 true라고 분류한것 중에서 실제 true인 것의 비율

recall: 실제 true인 것 중에서 모델이 true라고 예측한 것의 비율

 

저자는 이러한 문제를 해결하기 위해서 Soft Teacher를 제안했다.

soft teacher를 그림으로서 설명하면 아래와 같다.

soft teacher 설명

그림을 보면 teacher 모델이 생성한 빨간 박스와 student 모델이 생성한 녹색 박스가 있을때, 만약 두개의 IoU가 0.9 이하라고 하면 이를 background 라고 봐야하나 할때 약간의 애매함이 있을 수 있다. 저자는 이러한 부분을 완화시키고자 다음의 방법을 제시했다.

reliability 생성 흐름

student 모델이 생성한 녹색 후보 박스들을 다시 soft teacher 모델에 넣어서 background score를 계산한다. 그리고 이를 reliability score로 사용하여 background loss 계산 시 사용되는 가중치(w)를 계산할때 사용하는 방법이다.

수식은 아래와 같다.

soft teacher 수식

 

Box Jittering

Motivation

오른쪽 그래프는 저자가 실험한 내용으로 b를 보면, localization accuracy(IoU)와 후보 박스의 foreground score는 강한 양의 상관관계를 보여주지 않는다. foreground score가 1에 가까운데 반대로 IoU가 0에 가까운 경우도 있다. 그래서 저자는 score에 따라서 teacher generated pseudo box를 성택하는 것은 box regression에 적합하지 않고 더 나은 기준이 필요하다고 주장하였다.

그래프 c를 보면 저자가 찾은 다른 기준으로 box regression variance이다. 저자는 box regression variance가 낮을 수록 IoU가 높다는 것을 발견하였고, 이러한 variance를 계산하는데 jittering을 사용하였다.

그리고 모든 후보 box의 variance를 계산하는 것은 계산 비용이 너무 높으므로 foreground score가 0.5 보다 큰 box에 대한 reliability만 계산하였다.

 

 

box jittering을 그림으로 설명하면 다음과 같다.

bad/good regression

 

먼저 저자는 jitter라는 표현을 사용했는데 이는 여러개의 랜덤 오프셋을 주어서 box를 흐트러트린다는 의미이다.

위쪽 그림을 봤을때, teacher 모델이 생성한, 빨간 박스 bi 와 파란 박스 bi 가 있다. 빨간 박스의 경우 regression이 잘 된 경우이고 파란 박스는 bad regression인 경우이다.

먼저 빨간 박스(good regression인 경우)를 보면, box를 jitter 시켜서 생긴 jitterd box를 다시 teacher 모델에 넣는다면 어차피 object 근처에 box가 위치하므로 다시 원래 자리로 찾아 돌아가게 될것이고 그러므로 variance 또한 낮을것이라고 했다. 

반면 파란 박스(bad regression인 경우)를 보면, jittered 이미지를 다시 teacher 모델에 넣으면 box가 더 배경과 가까워지게 되고 완전히 흐터져버리게 될것이고 그러므로 variance 또한 높아질것이라고 한다.

그래서 이렇게 구한 variance를 이용하여 후보 박스에 대한 filtering을 적용하여 pseudo box를 생성한다.

 

저자의 경우 실험을 해본 결과 

 

box regression variance의 threshold의 경우 0.02에서 가장 성능이 좋았고, jittered box의 개수(N)의 경우 10개 이상 늘려도 성능에 유의미한 차이가 없다고 한다.

 

 

Experiments

다음은 저자의 실험 내용 중 일부이다.

 

Partially Labeled Data

먼저 이전 SOTA인 STAC와 본 논문의 성능을 비교한 실험이다. 이 실험에서는 COCO 데이터 중 train2017 데이터의 일부를 이용하여 각가 1%, 5%, 10% 의 labeled 데이터를 학습에 사용하였다. 본 논문(Ours)의 성능이 더 우수하였다.

 

Fully Labeled Data

다음은 detector와 backbone를 변경하면서 진행한 실험이다. train2017 데이터를 label 데이터로 모두 사용하였고, 추가로 unlabeled2017 데이터를 사용하여 성능 비교 실험을 진행하였다.

HTC++와 Swin-L을 사용한 경우 가장 성능이 높았다.