[CoIn] YOLOX Explanation — SimOTA For Dynamic Label Assignment

https://gmongaras.medium.com/yolox-explanation-simota-for-dynamic-label-assignment-8fa5ae397f76

YOLOX Explanation — SimOTA For Dynamic Label Assignment

위 글의 한국어 번역입니다.

 

Label Assignment in Object Detection

 

객체 탐지에서 레이블 할당(label assignment)은 학습 중에 어떤 bbox가 어떤 정답(ground truth, 이하 GT) 객체에 대응하는지를 결정하는 매우 중요한 작업입니다. 앞선 글에서 말했듯, 레이블 할당은 앵커들을 양성과 음성 집단으로 나눕니다.

 

양성 집단 앵커는 객체를 제대로 둘러싼 좋은 예측으로 간주되는 반면, 음성 집단 앵커는 객체를 둘러싸지 못한 나쁜 예측으로 간주됩니다. 예를 들어 아래의 이미지를 보겠습니다:

Bear Box

 

A, B, C로 라벨된 세 개의 bbox가 있다는 점이 있을 때, 만약 사람이 이 바운딩 박스를 양성과 음성으로 라벨링한다면, 아마 B는 곰의 머리를 완전히 둘렀으므로 양성, A와 C는 곰의 머리를 둘러싸지 못했으므로 음성이라고 할 것입니다. 곰의 머리가 정답 객체라면, B는 그 정답 객체와 매칭되고, A와 C는 어떤 정답과도 매칭되지 않거나, 배경(background)로 간주될 수 있습니다.

 

bbox를 양성/음성으로 라벨링하는 방법은 많으며, 이 글에서는 YOLOX가 SimOTA를 사용해 이 문제를 어떻게 푸는지를 설명합니다.

 

Why Use Label Assignment During Training?

 

SimOTA는 레이블을 할당하기 위해 GT가 필요합니다. 따라서 이는 학습 중에만 사용되며, 추론 시에는 사용되지 않습니다.

 

레이블 할당은 학습 안정성에 도움을 줄 수있습니다. 예를 들어, YOLOX에서 입력 이미지가 256일 때 예측 수가 약 1344개라고 해보면, 모든 예측을 최적화하는 대신, 레이블 할당을 통해 가장 좋은 예측들을 골라낸 다음, 그 좋은 예측들만 더 좋게 최적화할 수 있습니다. 참고로 음성 예측들이 학습에 아에 안 쓰이는 것은 아닙니다. 모든 양성 예측에 대해 회귀와 클래스를 최적화하고, objectness는 양성/음성 모두에 대해 최적화합니다.

 

이렇게하면 모델은 이미 좋은 예측을 더 좋게 만드는 데 집중하고, 나머지 예측에 대해선 크게 신경 쓰지 않습니다. 나쁜 예측을 가지치기 하면, 최종 모델의 예측 품질이 크게 향상됩니다.

 

학습이 더 안정적인 이유는, 모델이 처리해야할 경사 업데이트(gradient updates) 수가 줄어들기 때문입니다. 수천 개의 예측에 대해 회귀/클래스 손실을 업데이트하는 대신, 이미지마다 적은 수의 출력만 다루면 되므로, 최적화 공간이 더 다루기 쉬워집니다.

 

Label Assignment Before SimOTA

 

레이블 할당에는 여러 방식이 있습니다. 가장 흔한 방법 중 하나는, 정답 박스와 모든 예측 박스 사이의 IoU를 계산해 가장 높은 IoU를 가진 예측을 그 정답에 할당하는 방식입니다. 다른 알고리즘들도 유사한 방법을 사용하지만, SimOTA 저자들은 "각 정답 GT에 대해 context없이 독립적으로 양성/음성을 할당하는 것은 sub-optimal할 수 있다. 마치 문맥이 부족하면 부적절한 예측이 나오는 것과 같다." 라고 주장합니다.

 

이 sub-optimal label assignment 문제를 해결하기 위해, 저자들은 예전의 local context가 아닌 global context를 활용해 레이블을 할당할 것을 제안합니다.

 

OTA

 

OTA는 global context label assignment 문제를 다루기 위해 제안된 방법입니다. OTA는 레이블 할당 문제를 최적 수송(Optimal Transport, OT) 문제로 봅니다. OTA 저자들은 OT 문제를 다음과 같이 정의합니다:

 

"특정 지역에 m개의 공급자와 n명의 수요자가 있다. i-th 공급자는 s_i만큼의 물품을 갖고, j-th 수요자는 d_j만큼의 물품을 필요로 한다. 공급자 i에서 수요자 j로 물품 1단위를 옮기는 비용은 c_ij로 표시한다. 목표는 모든 물품을 최소 운송 비용으로 전달할 수 있는 운송 계획 π를 찾는 것이다."

 

요컨대, OT 문제는 공급자와 수요자가 있고, 최소 비용으로 운송하는 최적 계획을 찾는 것입니다.

 

OTA는 레이블 할당을 OT 문제로 공식화합니다. 여기서 m개의 공급자는 GT target이고, n명의 수요자는 이미지 위의 예측(앵커 위치) 입니다. 정답들은 양성 라벨은 수요자인 앵커들에 공급하고, 목표는 양성 라벨을 앵커에 최적으로 공금하는 계획을 찾는 것입니다. 즉, OT 문제를 통해 앵커/예측에 양성/음성을 가장 최적으로 라벨링하는 것이 목표입니다.

 

정답들과 함께, 또 다른 공급자가 있는데, 그것이 배경(background)입니다. 이 공급자는 나머지 라벨을 담당하며, 알고리즘의 4단계에서 나타납니다.

 

표기(Notation)

• I: 입력 이미지
• A: 입력 이미지 위의 앵커 집합
• G: 이미지 I의 정답 바운딩 박스들
• m: 정답 타깃의 개수
• n: 앵커의 개수
• k: 각 정답(gt)이 공급할 수 있는 양성 라벨 수
• sᵢ: i번째 정답의 공급량 (sᵢ = k)
• dⱼ: j번째 앵커의 수요량
• c: 정답 gtᵢ에서 앵커 aⱼ로, 양성 라벨 1개를 보내는 비용
• Ø: 배경 클래스(보통 0으로 표기)
• α: 회귀 손실 가중 계수 (일반적으로 ≥ 1, 회귀 손실을 더 크게 가중)
• T: Sinkhorn-Knopp 반복 횟수

 

알고리즘은 다음과 같습니다:

https://arxiv.org/abs/2103.14259

 

OTA 알고리즘 개요:

1. m, n을 정답 수와 앵커 수로 설정
2. 이미지 I를 모델에 통과시켜 클래스 예측 Pᶜˡˢ와 회귀 예측 Pᵇᵒˣ을 얻음
3. 길이 m+1의 공급 벡터 s를 만들고, 동적 k 추정(dynamic k)으로  각 정답의 공급량을 채움
4. s[m+1] = n − sum(s)
   → 공급 벡터의 m+1 위치(배경)의 공급량은 n − sum(s)
5. 길이 n의 수요 벡터 d를 1로 초기화
6. 각 예측 j와 정답 i의 쌍에 대해 클래스 쌍별 손실:
   cᶜˡˢ = FocalLoss(Pᶜˡˢ, Gᶜˡˢ)
7. 각 예측 j와 정답 i의 쌍에 대해 회귀 쌍별 손실:
   cʳᵉᵍ = IoULoss(Pᵇᵒˣ, Gᵇᵒˣ)
8. 각 앵커 j와 정답 i의 쌍에 대해 센터 프라이어(center prior):
   cᶜᵖ = CenterPrior(Aⱼ, Gᵇᵒˣ)
9. 배경 비용: cᵇᵍ = FocalLoss(Pᶜˡˢ, Ø)
10. 전경 비용: cᶠᵍ = cᶜˡˢ + α·cʳᵉᵍ + cᶜᵖ
11. 최종 비용 행렬 c를 cᵇᵍ를 cᶠᵍ에 이어 붙여 (m+1, n) 모양으로 구성
12. u, v를 1로 초기화
13. Sinkhorn-Knopp 반복을 T번 수행해 u, v를 채움
14. 원 논문 식에 따라 최적 할당 계획 𝝅* 계산
15. return 𝝅*

SimOTA

 

겉보기에 복잡해 보이지만, 실제로는 그렇지 않습니다. 한편, SimOTA는 이 알고리즘을 훨씬 간단하고 빠르게 만든 알고리즘입니다. YOLOX의 저자들은 OTA가 성능을 높이긴 하지만, 학습 속도를 약 25% 느리게 만든다는 점을 발견했습니다. 학습을 수백 iteration단위라고 생각한다면, 이는 상당한 비용입니다. 따라서 저자들은 Sinkhorn 반복을 제거하고 최적 할당 계획을 근사함으로써, 성능 향상을 유지하면서도 속도를 크게 높일 수 있음을 보였습니다.

 

SimOTA는 Sinkhorn을 쓰는 대신, 각 정답 i에 대해 비용이 가장 낮은 상위 k_i개의 예측을 양성으로 선택합니다. 이 방식 덕분에 최적화 알고리즘으로 가장 최적의 할당을 찾는 대신, 정답들을 한 번만 순회하여 할당을 근사할 수 있습니다.

 

SimOTA 알고리즘 개요:

 

• m, n을 정답 수와 앵커 수로 설정
• 이미지 I를 모델에 통과시켜 클래스 예측 Pᶜˡˢ와 회귀 예측 Pᵇᵒˣ을 얻음
• 길이 m+1의 공급 벡터 s를 만들고, 동적 k 추정으로 각 정답의 공급량을 채움
• s[m+1] = n − sum(s)
  → 공급 벡터의 m+1 위치(배경)의 공급량은 n − sum(s)
• 길이 n의 수요 벡터 d를 1로 초기화
• 각 예측 j와 정답 i의 쌍에 대해 클래스 쌍별 손실:
  cᶜˡˢ = FocalLoss(Pᶜˡˢ, Gᶜˡˢ)
• 각 예측 j와 정답 i의 쌍에 대해 회귀 쌍별 손실:
  cʳᵉᵍ = IoULoss(Pᵇᵒˣ, Gᵇᵒˣ)
• 각 앵커 j와 정답 i의 쌍에 대해 센터 프라이어:
  cᶜᵖ = CenterPrior(Aⱼ, Gᵇᵒˣ)
• 배경 비용: cᵇᵍ = FocalLoss(Pᶜˡˢ, Ø)
• 전경 비용: cᶠᵍ = cᶜˡˢ + α·cʳᵉᵍ + cᶜᵖ
• 최종 비용 행렬 c를 cᵇᵍ를 cᶠᵍ에 이어 붙여 (m+1, n) 모양으로 구성
• 각 정답 i에 대해, 공급량 sᵢ만큼 비용이 가장 낮은 상위 sᵢ개의 예측을 선택
  → 결과 배열은 길이 m이며, 각 인덱스 i에는 최대 sᵢ개의 예측이 담긴다.
• 결과 배열을 반환

 

SimOTA의 수행 결과는 길이 m의 배열이며, 배열의 i-th 원소는 정답 G_i에 대응하는 양성 앵커/예측들입니다. 이 배열에 포함되지 않은 나머지 예측들은 정답 매칭 없이 음성으로 간주됩니다.

 

Dynamic K Estimation

 

k는 각 정답 객체의 공급량이며, 계산 방식은 두 가지가 있습니다. 단순한 방식은 모든 정답에 대해 k를 상수로 두는 것입니다. 하지만 이렇게 하면 모든 정답에 동일한 수의 앵커를 할당하게 되는데, 이는 바람직하지 않습니다.

 

두 번째 방법은 각 정답에 대해 개별적으로 k를 할당하는 것입니다. OTA의 저자들은 Dynamic k estimation을 제안하며, 각 정답의 공급량을 근사합니다. 방법은 다음과 같습니다:

 

"모든 예측을 살펴 각 예측과 정답 사이의 IoU를 기준으로 상위 q개의 예측을 고른 뒤, 이 상위 q개의 IoU 값을 합산해 그 정답의 k값으로 사용한다."

 

이 방법을 사용하면, 각 정답에 대해 양성 라벨 수(공급량)를 그 정답을 "얼마나 정확히 둘러싸는 예측이 많은지"로 추정할 수 있게 됩니다. 그 결과, SimOTA를 사용할 때 해당 정답을 잘 회귀하는 앵커가 많을수록, 그 정답에 더 많은 양성 앵커가 할당되는 경향을 가집니다. 한편, k는 더 이상 사람이 바꿔야 할 하이퍼-파라미터가 아니지만, 이제 q가 튜닝해야 할 파라미터가 됩니다. 논문에서는 q = 20을 사용했으며 무난하게 잘 동작했다고 합니다. 아래는 dynamic k estimation에 대한 작성자 코드입니다:

# The supplying vector
s_i = np.ones(m+1, dtype=np.int16)
# The sum of all k values
k_sum = 0
# Iterate over all ground truth boxes (i = gt_i)
for i in range(0, m):
    # Get the ith truth value
    gt = G_reg[i]
    # Get the (x, y) coordinates of the intersections
    xA = np.maximum(gt[0], P_reg[:, 0])
    yA = np.maximum(gt[1], P_reg[:, 1])
    xB = np.minimum(gt[0]+gt[2], P_reg[:, 0]+P_reg[:, 2])
    yB = np.minimum(gt[1]+gt[3], P_reg[:, 1]+P_reg[:, 3])
    # Get the area of the intersections
    intersectionArea = np.maximum(0, xB - xA + 1) * np.maximum(0, yB - yA + 1)
    # Compute the area of both rectangles
    areaA = (gt[2]+1)*(gt[3]+1)
    areaB = (P_reg[:, 2]+1)*(P_reg[:, 3]+1)
    # Get the union of the rectangles
    union = areaA + areaB - intersectionArea
    # Compute the intersection over union for all anchors
    IoU = intersectionArea/union
    # Get the q top IoU values (the top q predictions)
    # and sum them up to get the k for this gt
    k = np.sort(IoU)[-q:].sum()
    # Add the k value to the total k sum
    k_sum += k
    # Save the k value to the supplying vector
    # as an iteger
    s_i[i] = int(round(k))

 

Center Prior

 

SimOTA가 레이블을 할당할 때 부딪히는 문제 중 하나는 하나의 GT가 갖는 양성 라벨이 아무 앵커 예측에나 배정될 수 있다는 점입니다. 가끔은 어떤 GT도 양성 라벨을 줄 만한 좋은 후보를 충분히 갖지 못할 수 있는데, 그대로 k개의 양성 라벨을 반드시 배정해야 하므로, 결국 질 나쁜 양성 라벨을 자신에게 배정하게 되는 것입니다.

 

예를 들어 아래의 이미지를 보면:

 

 

이미지의 모든 bbox가 곰의 얼굴을 잘 덮지 못합니다. 곰의 얼굴이 GT객체이고 이 GT의 공급량(=k값)이 2라고 하면, 이 GT는 두 개의 예측을 양성 라벨로 뽑아야 합니다. 분명히 좋은 예측 두 개가 없으므로, 이 GT 객체는 결국 형편없는 두 예측을 양성으로 만들게 됩니다.

 

학습 초기 단계에서는 이런 나쁜 예측 문제가 특히 두드러집니다. 따라서 GT가 고를 만한 좋은 예측이 부족한 상황은 큰 문제입니다. 해결책은 반지름 r을 정하고, 각 GT 중심에서 가장 가까운 r^2개의 앵커는 패널티 없이 두는 한편, 그 r^2 반경 밖의 모든 앵커에는 추가 비용(extra cost)을 더해 패널티를 주는 것입니다.

 

SimOTA 알고리즘 (8)단계에서 보듯, 이렇게 계산한 center prior 비용을 모든 예측의 총 비용에 더합니다. 왜 r^2 반경 안의 예측들이 해당 GT의 양성 앵커로 선택될 가능성이 더 높도록 설계했을까요? 직관은 이렇습니다:

 

"GT와 가까운 앵커일수록 최적화가 쉽고, 먼 예측일수록 최적화가 어렵다."

 

GT의 중심 G는 x좌표에 너비의 절반과 y좌표에 높이의 절반을 더해 구할 수 있습니다:

center = (G[0] + G[2]//2,  G[1] + G[3]//2)

 

그 다음, 모든 앵커 A와 GT의 중심 G 사이의 거리를 다음처럼 계산합니다:

diff = A - center
dist = np.sqrt(np.sum(diff**2, axis=-1))

 

아래는 단일 GT와 모든 앵커 사이의 center prior를 계산하는 코드입니다:

"""
Get the center prior between a gt and the anchor locations
on the image
Inputs:
  A - All anchors for a single image
  G_box - A ground truth box for this image
  r - radius used to select anchors in this function
  extraCost - The extra cost to add to those anchors not in
              the r**2 radius
Output:
  Array with the same number of values as the number of anchors
  where each value is the center prior value of that anchor
"""

def centerPrior(A, G_box, r, extraCost):
    """
    Center Prior selects the r**2 closest anchors according to the
    center distance between the anchors and gts. Those anchors
    that are in the radius are not subject to any extra cost, but those
    anchors outside the radius are assigned extra cost to avoid
    having them be labelled as positive anchors for this gt.
    """
    
    # Get the center location of the ground truth boudning box
    center = (G_box[0]+(G_box[2]//2), G_box[1]+(G_box[3]//2))
    
    # Get the difference between the center locations in A and the
    # center location of the gt bounding box
    diff = A-center
    
    # Use the distance formula to get the distance from the
    # gt center location for each anchor
    dist = np.sqrt(np.sum(diff**2, axis=-1))
    
    # Get the indices of the distances which are greater
    # than r**2 meaning the anchor is outside the radius
    idx_neg = np.where(dist > r**2)
    
    # Array of zeros corresponding to the center prior of each anchor
    c_cp = np.zeros(A.shape[0])
    
    # Those values not in the r**2 radius are subject to a constant cost
    c_cp[idx_neg] = extraCost
    
    return c_cp

 

각 FPN의 레벨이 서로 다른 앵커 수를 가질 수 있기 떄문에, 입력 이미지가 256 픽셀일 때, stride가 32, 16, 8인 FPN 레벨의 앵커 수는 64, 256, 1024입니다. 만약 반지름 r이 꽤 크고 곰의 머리가 GT라면, 아래의 그림에서 파란 원 안의 그리드 교차점들이 추가 비용을 부과하지 않는 앵커 포인트가 됩니다:

 

OTA의 저자들은 r^2에 가장 가까운 앵커를 고르는 것이 학습을 (특히 학습 초기 단계에서) 안정화한다고 말합니다. 또한, 이런 추가적인 안정성 덕분에 모델의 성능이 더 좋아진다고 주장합니다.

 

YOLOX는 실제로 r^2 반경 밖의 라벨들을 가지치기(prune)합니다. 즉, 반경 밖 예측에 추가 비용을 주는 대신, r^2 밖의 모든 예측들을 제거합니다. 때로는 어떤 GT에 대해 예측이 하나도 남지 않는 경우도 생기지만, 반경 밖 예측을 제거하면 center prior를 유지하고 몇 개의 나쁜 예측 때문에 gradient 폭주가 일어나는 것을 막아 학습 안정성에 도움을 줍니다.

 

다음 글에서는 YOLOX가 사용하는 데이터 증강(data augmentation)을 설명하겠습니다.