YOLO v1 논문 리뷰 및 코드 구현

김상현·2021년 4월 23일

논문 제목: You Only Look Once(YOLO): Unified, Real-Time Object Detection

YOLO v1 개요

R-CNN과 같은 이전의 object detection 모델과 다르게 YOLO는 1-stage detector로 하나의 네트워크로 detection을 수행한다. 하나의 네트워크를 사용하므로 end-to-end 학습이 가능하다.

YOLO의 이점은 다음과 같다.

빠르다.
- 초당 45프레임을 처리할 수 있고, fast version에서는 초당 150프레임을 처리할 수 있다.
- 따라서 real-time detection이 가능하며, 기존의 다른 real-time detector들 보다 2배 이상의 mAP를 보여준다.
추론 시 이미지를 전역적으로 파악한다.
- 기존의 sliding window와 region proposal과 다르게 이미지의 전체를 본다.
- 따라서 해당 클래스의 맥락적 정보 뿐만 아니라 모습(appearance)을 encode해서 fast R-CNN보다 background error를 절반 가량 줄일 수 있다.
객체의 일반화 가능한 표현(representation)을 익힌다.
- 자연 이미지(natural image)로 학습한 후 그림(artwork)에서 test를 진행해도 다른 모델들 보다 좋은 성능을 보인다.

물론 당시 state-of-the-art detection system들 보다 정확도 측면에서 뒤떨어지긴 했지만 위의 3가지 장점을 갖는 모델이다.

Unified Detection

사진1

YOLO는 입력 이미지를 SxS grid로 나눈다. 만약 개체의 중심이 grid cell에 들어간다면, 그 grid cell은 해당 객체에 대해 detecting을 수행한다.

각각의 grid cell은 B개의 bounding box와 confidence score를 예측한다. confidence score는 박스가 물체를 포함하는지와 얼마나 정확하게 예측한 박스가 맞는지를 반영한다. confidence는 다음과 같이 정의된다.

confidence=Pr(object)*IOU(pred, truth)

즉, 해당 박스에 물체가 존재하지 않을 때는 0이 되고, 물체가 존재할 때에는 ground truth와 prediction box의 IOU(Intersection over Union)값을 갖게 된다.

각각의 bounding box는 $x,y,w,h,confidence$ 를 예측한다. $x,y$ 는 중심 좌표를 나타내고, $w,h$ 는 높이와 너비를 나타낸다. 마지막으로 confidence는 위에 정의한 confidence이다.

각각의 grid cell은 C개의 conditional class probabilities $Pr(class_i|object)$ 를 예측한다. bounding box의 개수와 관계없이 하나의 grid cell에서 하나의 class probabilities를 예측한다. 테스트 시에 다음과 같은 식으로 class-specific confidence를 얻는다.

Pr(Class_i|object)*Pr(object)*IoU(pred, truth) = Pr(class_i)*IoU(pred, truth)

class-specific confidence는 class가 box에 등장할 확률과 객체가 예측된 box에 얼마나 잘 맞는지를 나타낸다.

본 논문에서는 PASCAL VOC dataset에 맞춰 S=7, B=2, C=20을 사용했다. 따라서 모델의 output은 7X7X30 tensor가 된다.
cf)실제로는 배치를 고려한 output이 나온다. 또한 30 = (20 + 2*5)이다.

사진2. YOLO output shape

위의 사진은 확인하면 각각의 grid cell은 1x30의 벡터인 것을 확인할 수 있다. 각각의 bounding box는 위의 설명과 같이 5개의 값을 예측한다. 따라서 앞의 10개 값은 각각의 bounding box의 예측을 나타낸다. 그 뒤의 20개의 값은 VOC dataset이 20개의 class를 갖기 때문에 20개의 conditional class probabilities를 예측한다.

2.1 Network Design

사진3. Model Architecture

네트워크 구조는 위의 그림과 같다. Image classification을 위한 GoogLeNet 구조에 영감을 받아 해당 네트워크 구조를 설계했다. 기존의GoogLeNet은 Inception module을 사용한 반면에 YOLO에서는 Inception module을 일자로 이어둔 모델을 사용했다.
Fast version YOLO의 경우 네트워크에서 convolutional layer의 수를 24개에서 9개로 줄이고, 적은 filter를 사용했다. 네트워크 구조 이외에 모든 parameter는 original YOLO와 같다.
네트워크의 출력은 7x7x30 tensor 이다. (VOC dataset에 맞춘 output 크기)

2.2 Training

YOLO는 본 학습을 시작하기 전 GoogLeNet을 변형시킨 앞의 20개의 convolutional layers를 ImageNet dataset을 이용해 classification을 수행하는 방법으로 pre-train 한다. detection을 수행하기 위한 network로 바꾸기 위해 pre-train된 network에 4개의 convolutional layers와 2개의 fully-connected layers를 추가한다. pre-train시 224x224 image를 input으로 사용했으나, detection시 object detection 특성상 fine-grained visual information을 요구하므로 448x448 image의 input을 사용한다.

네트워크 출력인 height와 width를 normalize를 통해 0~1사이로 유계시킨다. 또한 bounding box의 좌표인 $x,y$ 또한 parameterize를 통해 0~1로 유계시킨다.

활성화 함수로 LeakyReLU(0.1)을 사용했다.
최적화를 쉽게 하기 위해 sum-squared error를 사용한다. 하지만 다음과 같은 이유로 모델의 불안정성을 발생시킨다.

localization error와 classification error를 동등하게 가중치를 주는 것은 비이상적일 수 있다.
객체가 존재하지 않는 grid cell이 많기 때문에 객체가 존재하는 cell의 confidence가 0을 향하게 하는 영향(gradient)을 미칠 수 있다.

이러한 문제를 해결하기 위해 bounding box coordinate loss를 증가시키고, 객체가 존재하지 않는 box의 경우 confidence loss를 감소시켰다. 이를 위해 $\lambda_{coord}=5$ , $\lambda_{noobj}=0.5$ 를 사용했다.
또한 sum-squared error는 box 크기에 상관없이 같은 error를 사용해서 이런 문제를 해결하기 위해 bounding box의 width와 height을 제곱근을 이용해 계산한다.

YOLO는 여러 개의 bounding boxes를 각각의 grid cell에서 예측한다. 학습시에 각 grid cell마다 한 개의 bounding box를 원하기 때문에 bounding boxes 중에서 ground truth와 가장 IoU가 높은 box만 선택해서 학습을 시킨다. 이때 선택된 한 개의 bounding box를 "responsible"이라 할당한다.

Loss function은 다음과 같다.

사진3. loss function

$1^{obj}_{i}$ 는 cell i에 객체가 존재하는가를 나타내는 indicator function이다. $1^{obj}_{ij}$ 또한 indicator function으로 cell i에 j번째 box가 "responsible"인지 나타낸다. $C_i$ 는 cell i에서 confidence score를 나타낸다. 위의 loss function을 최적화하면서 학습을 진행한다.

논문의 저자들은 다음과 같은 방법으로 모델을 학습시켰다.

Epoch = 135, batch size = 64, momentum = 0.9, decay = 0.0005
learning rate scheculing: 첫 epoch에서 $10^{-3}$ 로 시작해서 75 epoch까지 $10^{-2}$ 으로 학습시킨다. 이후 30 epochs 동안 $10^{-3}$ 으로 학습하고, 마지막 30 epochs 동안 $10^{-4}$ 으로 학습시킨다.
overfitting을 막기 위해 dropout과 data augmentation을 활용했다.

2.3 Inference

학습을 마친 YOLO 모델은 PASCAL VOC의 이미지에 대해 각각 98개의 bounding boxes를 출력한다. 이렇게 나온 98개의 bounding boxes들에 대해 NMS(Nom-Maximum Suprression)을 적용한다.

2.4 Limitations of YOLO

각각의 grid cell에서 2개의 예측한 bounding boxes를 얻고, 하나의 class만 갖을 수 있으므로 근접한 작은 물체에 대해 잘 감지 못 한다.
Train dataset에 존재하는 bounding box의 형태와 다를 경우 일반화하는데 어려움이 있다. 또한 입력 이미지를 여러 번 down-sampling(여러 layers를 통과)해서 coarse features를 이용해 bounding box를 예측한다.
큰 상자에서 작은 error보다 작은 상자에서 작은 error를 더 크게 영향을 받아 localization의 부정확한 경우가 있다.

Experiments

사진4. 성능 비교

위의 표를 보면 YOLO가 꽤 준수한 mAP를 보여주면서 동시에 굉장히 빠른 처리 속도를 자랑하는 것을 확인할 수 있다.

사진5. Error Analysis

YOLO의 장점에서 얘기한 것처럼 background error가 다른 모델에 비해 적은 것을 알 수 있다. 이미지를 전역적으로 파악해서 위와 같은 이점을 얻을 수 있다고 논문의 저자들은 밝혔다.

Conclusion

YOLO는 unified real-time detector로 쉽게 만들 수 있고, 전체 이미지로 직접 학습되어진다. 이전의 모델들 보다 굉장히 빠른 속도를 보여줬고, 꽤 준수한 detection 성능을 보여줬다.

코드 구현

코드는 개인 github에 pytorch를 사용해 구현했다.
aladdinpersson님의 github와 코드 구현 영상(youtube)을 참고해 코드 구현을 진행했다.
코드 중간에 confidence score를 이용해 loss를 구하지 않는 것을 확인해 그 부분을 수정하고, 다른 부분들에 주석을 통해 모델의 작동방식을 설명했다.