[논문 리뷰] YOLOv6

개요

논문 링크: https://arxiv.org/abs/2209.02976

github 링크: https://github.com/meituan/YOLOv6

이전 글: YOLOv5 정리

제목: YOLOv6: A Single-Stage Object Detection Framework for Industrial Applications

노트:  YOLO v7보다 조금 늦게 나왔는데 v6라고 불린다. 성능대비 FPS가 기존의 YOLO보다 좋다. 

요약:

• EfficientRep backbone 제안
• Task alignment learning 활용 (label assignment 잘하는 방법)
• anchor free 방법 사용 (anchor box 가 없다는 뜻)
• VariFocal loss 도입
• 다양한 bbox regression loss 실험 (다양한 IoU loss, distributed focal loss)
• Objectness branch 제거
• Gray border padding 없이 성능 유지하는 방법 제안
• Self distillation 도입
• Post training quantization 잘되는 세팅 제안
• Quantization aware training 잘하는 방법 제안

방법

Network architecture

디테일은 접은글 참고. RepVGG block과 CSP block을 효율적으로 합친 CSPStackRep을 사용한다는 것이 핵심이다

그림들은 A Comprehensive Review of YOLO 에서 빌려왔다

우선 그림들에서 사용되는 블록들은 다음과 같다

Backbone

• YOLOv6에서 CSPStackRep이라는 블록을 제시하는데 BepC3StageBlock이랑 동일한 블록을 의미한다.
  • RepVGG block을 그대로 사용하지 않는 이유는 network size를 키울 때 계산 량과 parameter 수가 너무 많아지기 때문에
  • CSPStackRep는 CSP block과 RepVGG block을 섞은 block인데, 1x1 conv로 channel 수를 줄인 후 RepVGG block (RepConv)을 사용하기 때문에 효율이 좋아진다

• Backbone 구조는 다음과 같다 (논문에서 EfficientRep라고 부른다)

Neck

• YOLO-v4/v5처럼 PAN을 사용하지만 비슷하지만 CSPStackRep를 사용 (논문에서 RepPAN이라고 부른다)

Head

• decoupled head 사용 (classification, bbox head를 따로 분리한다는 뜻)
  • classification, bbox는 각각 3x3 conv + 1x1 conv로 이뤄진다. 3x3 conv를 1 layer만 사용한 것이 특징이라고 하기는 한다
• anchor point based regression paradigm을 사용하는데, anchor point에서부터 bbox 변들의 거리를 예측한다는 뜻이다

Architecture Summary

Label assignment

• 원래 YOLOX처럼 ([논문 리뷰] YOLOX) SimOTA를 사용했지만 TAL 사용하는 것이 더 좋다고 한다 (SimOTA을 사용하면 학습이 느리고 unstable 했다고 한다)

TAL (Task Alignment learning) 간단 설명

• 링크: https://arxiv.org/abs/2108.07755
• 논문에서 풀고자 하는 문제는 다음과 같다: detection문제에서 object classification과 bbox regression은 서로 다른 task이지만 같은 feature에서부터 예측이 이뤄진다. 어떻게 이 두 task를 잘 align 할 수 있을까? 몇 가지 제안을 하는데 그중 하나가 TAL이다
• 우선 \( t=s^\alpha u^\beta\) 점수 기반으로 label asignment를 한다
  • alpha=1, beta=6
  • s = classification score (GT class에 해당되는 예측 확률)
  • u = IoU (GT bbox와 IoU)
  • t는 classification과 bbox regression 중 중점을 어디에 줄지 조절할 수 있게 해준다
• 각 GT label마다 t 값이 가장 큰 m개를 positive sample로 사용하고, 나머지는 negative sample로 사용
  • m=13

Loss function

Classification

• VariFocal loss 사용 (https://arxiv.org/abs/2008.13367)

• p = predicted IACS
• q = IACS (GT class는 GT bbox와 predicted bbox의 IoU, 나머지는 0인 벡터)
• 여기서 IACS는 IoU Aware Classification Score의 약자다. 보통 classification score은 GT class에 1, 나머지는 0인 벡터로 표현이 된다면, IACS는 GT class에 GT bbox와 predicted bbox와의 IoU 값을 갖고, 나머지는 0이다.
• 특징: negative example (q=0) 같은 경우만 \( p^\gamma\)로 down-weight 하고, positive sample은 down-weight 하지 않는다. 참고로 Focal loss 같은 경우 positive sample에도 비슷하게 \( (1-p)^\gamma\)factor이 있다 

y=1: GT class, p = predicted probability

• VariFocal Loss 논문에서 \( \alpha=0.75, \gamma=2.0\) 사용

Box Regression

• IoU loss 같은 경우, GIoU 또는 SIoU loss를 사용 (N,T 모델은 SIoU, 나머지는 GIoU loss)
• GIoU: loss = 1-GIoU, GIoU는 다음과 같이 정의된다 (https://arxiv.org/abs/1902.09630)

• SIoU: 어렵지는 않지만 식이 좀 복잡하기 때문에 자세한 내용은 논문 참고 (https://arxiv.org/abs/2205.12740). Angle, distance, shape cost, IoU를 활용해서 정의한다
  
  • \( \Lambda\) = angle cost: GT & predicted bbox를 x 축 또는 y 축 위에 놓이도록 align 시키는 역할을 한다 (학습할 때 predicted bbox와 GT bbox 사이 각도가 왔다 갔다 하는 것을 막아주는 역할)
  • \( \Delta\) = distance cost: bbox 중심사이 거리를 작게 하는 역할을 한다. 이때 distance cost는 angle dependent 하게 정의된다 (angle cost 활용)
  • \( \Omega\) = shape cost: GT & predicted bbox의 width, height를 맞춰준다

SIoU loss

• Distriburted focal loss (DFL) (https://arxiv.org/abs/2006.04388)같은 경우 작은 모델에서 사용하면 inference speed 저하가 있기 때문에 M/L 모델에서만 사용. 보통 bbox 위치를 직접 regression 하지만, DFL를 활용하기 위해서 heat map 같은 approach를 사용한다.
  • 우선 \( \hat{y} = \sum_{i=0}^n p(y_i) y_i\)로 적는다. 여기서 \( y_i\)는 bbox regression branch에서 각 feature에서 예측되는 bbox 위치가 되고, \( P(y_i)\)는 각 feature에서 예측 값이 얼마나 믿을만한지 알려주는 척도 (확률분포)로 해석 가능하다.  이 확률 분포를 기반으로 예측 값들을 평군내서 최종 bbox 위치 예측 값 \( \hat{y}\)을 구한다. (밑에 그림에서 DFL 참고)
  • \( P(y) \)는 \( S\)라는 별도의 softmax layer를 통해서 구한다. 편의상 \(S_i = P(y_i)\)라고 부른다
  • \(DFL(S_i,S_{i+1}) = -((y_{i+1}-y)\log S_i + (y-y_i) \log S_{i+1})\)
  • DFL은 y 근처에 확률이 높게 나오도록 한다. Global minimum은 \( S_i = \frac{y_{i+1}-y}{y_{i+1}-y_{i}}\),  \( S_{i+1} = \frac{y-y_{i}}{y_{i+1}-y_{i}}\)이고, 이때 \( \hat{y} =\sum_{i=0}^n p(y_i) y_i= y\)인 것을 보일 수 있다.

Objectness

• Objectness loss를 사용하면 성능이 저하되어서 사용하지 않는다
• Objectness branch를 추가하면 task alignment learning에 방해가 되어서 성능 저하가 일어난다고 저자들은 추측한다

Training

Epoch

• YOLOv5 같은 경우 300 epoch 학습했다
• YOLOv6에서는 오래 학습하면 도움이 되어서 400 epoch동안 학습을 한다

Self-Distillation

• self-distillation이기 때문에 teacher과 student의 모델 구조는 같다
• 우선 pre-training을 하고, pre-trained 모델을 teacher로 사용한다
• DFL를 사용하면 regression도 classification task로 바뀌기 때문에 (DFL 같은 경우 확률 분포를 내밷는다), classification과 bbox regression 둘 다 knowledge distillation이 가능하다
• \( L_{KD} = KL(p^{cls}_{t} || p^{cls}_s) + KL(p^{reg}_t || p^{reg}_s)\)
  • p = 확률, cls = classification prediction, reg = regression prediction, t = teacher, s = student
• \( L_{total} = L_{det} + \alpha L_{KD}\)
  • det = detection (이 전 section에서 다룬 loss들)
  • alpha: KD loss에 주는 weight; 처음에는 높게 잡지만 학습 끝날쯤 0이 되도록 schedule 해서 학습 (cosine weight decay)

Gray border of images

• half-stride gray border을 이미지에 추가하면 이미지 = border에 있는 물체를 detection 할 때 도움이 된다고 한다 (YOLOv5, v7에서 이 트릭을 사용)
• 하지만 이것은 inference speed를 저하시킨다
• gray border을 추가해야 성능이 좋아지는 이유는 mosaic augmentation 때문이라고 저자들은 추측한다 (Mosaic augmentation에서 gray border padding을 사용하기 때문)
• 학습이 끝날 때쯤 mosaic augmentation을 사용하지 않는 trick과 gray border의 영역을 바꾸는 trick을 사용하면 gray border 없이도 잘 작동한다고 한다

Quantization

• RepOptimizer로 학습 (RepVGG block을 그냥 학습하면 post training quantization이 잘 안 됨, [간단 리뷰] RepOptimizer (RepOpt-VGG))
• Quantization이 잘 안 되는 layer들은 floating point로 유지하는 trick 사용
• QAT 사용; 이때 channel-wise distillation도 활용 (teacher = float32 모델, student = int8 quantized 모델, 그림 참고)

학습 세팅

• Optimizer: SGD with momentum
• learning rate: cosine decay
• Warm-up, grouped weight decay, exponential moving average
• Mosaic, mixup augmentation
• hyperparameters들을 github 참고