YOLO - dataset load, image size - resize, compress

미래내일 일경험 프로젝트를 진행하면서 2TB 정도의 대용량 데이터셋을 이용하면서 AI를 학습할 때 힘들었던 부분이 두 가지 있었다.

1. 데이터셋 저장

2. 데이터셋 load

데이터 저장과 데이터셋 로드는 밀접한 연관이 있다. 데이터셋을 어디에 저장하냐에 따라 AI 학습시 데이터를 load하는 속도에 영향으 주기 때문이다.

 

1. Dataset

1. Google Drive + Colab

처음 학습은 Google Drive와 Colab을 이용했다. 2TB의 데이터셋은 Google Drive에 저장을 했으며, AI 학습은 Colab에서 진행했다. 하지만 Colab은 무료로 사용하면 무료로 제공되는 디스크의 크기는 80GB 정도이며 Pro의 경우 200GB 정도를 사용할 수 있다. 따라서 Google Drive에 데이터셋을 저장하고 load하는 방법으로 AI 학습을 해야 했다. 하지만 Google Drive는 초당 불러올 수 있는 데이터 수에 제한이 있어 2TB를 가져오는 것은 거의 불가능에 가까웠다.

train: Scanning /content/drive/MyDrive/food_working/train/labels/01011001... 
6215 images, 0 backgrounds, 175 corrupt:   1% 6390/619850 [2:03:14<232:47:10,  1.37s/it]

YOLO를 이용해 학습을 진행하면 데이터를 Scanning하는 작업을 하는데 619,850장의 이미지를 불러오면서 Scanning 작업을 한다. 하지만 위의 결과를 보면 알 수 있듯 6,390장을 불러오는데 2시간이 넘게 걸렸다. 따라서 Colab을 활용한 AI 학습은 포기했다.

 

2. 엘리스(Elice)

Goole Drive와 Colab을 이용한 학습을 포기하고 어쩔 수 없이 새로운 학습 방법을 찾아야 했다. 여기저기 뒤지면서 찾다가 LGAimers에 참가했을 때 Elice에서 제공한 클라우드 환경에서 GPU(무려 A100)를 이용하면서 해커톤을 했던 것이 기억났다. 그래서 Elice를 이용해 AI 학습을 진행했다. Elice의 경우 드라이브와 GPU를 선택해 학습을 할 수 있다. 드라이브의 용량이 클 수록, GPU의 성능이 좋을 수록 가격이 비싸다. 필자의 경우 1TB + A100 80GB를 이용해 시간당 2,000원 정도의 비용이 나왔다. 

Class    Images   Instances    Box(P        R        mAP50      mAP50-95
all      35233      70466      0.921      0.912      0.956      0.854

위의 결과 처럼 학습을 완료하고 좋은 결과를 얻을 수 있었다.

 

2. Image Size

분명 2TB의 데이터셋을 이용했다고 말했는데 엘리스를 이용해서 학습을 할 때 드라이브를 1TB로 사용했다고 말했다. 데이터셋의 크기를 줄였다. 물론 객체 탐지 갯수를 줄이기도 했지만 이미지 사이즈 하나하나의 용량을 줄였다. 두 가지 방법을 사용했다.

1. resize - 학습에 맞는 이미지 크기(640)

이 방법은 데이터셋의 특성에 따라 어려울 수도 있고, 쉬울 수도 있다. 객체 탐지를 위해서는 좌표 정보가 필요하다. 즉, annotation: 객체의 위치와 범주에 대한 정보가 있는데 YOLO의 경우 Class id, 중심 x좌표, 중심 y 좌표, 너비, 높이에 대한 정보가 필요하다. 하지만 이미지의 비율을 변경하면 좌표가 바뀌기 때문에 좌표에 대한 수정도 필요하다. 예를 들어 아래와 같다.

사진1

사진1은 YOLO 학습에 맞게 640x640으로 비율을 조정하고 annotation 정보는 수정하지 않고 학습했을 때 train 결과다. 필자가 학습에 활용한 음식 사진의 경우 대부분 정면에서 음식 하나를 가운데에 놓고 찍어 놓은 사진이다. 따라서 비율을 조정해도 음식은 그대로 가운데에 있어서 좌표의 영향을 받지 않았다. 

하지만 대부분의 경우 비율에 맞게 annotation에 있는 좌표 정보도 수정해야 한다.

사진2

사진2를 보면 왼쪽이 train 결과이고 오른쪽이 val 결과이다. YOLO-pose에 대한 학습이지만 객체 탐지보다 더 많은 좌표를 사용하기에 좋은 예시라 판단했다. 왼쪽 사진의 경우 학습할 때부터 이미 keypoint 좌표가 사람 모양에서 다르다. 따라서 val 결과 역시 keypoint의 위치가 다르다. 당연히 mAP50, mAP50-95의 성능 역시 굉장히 좋지 않다. 아래와 같다. 

original image: 0.99006	0.80327
resize image: 0.95113	0.57925

GPU	epochs	batch size	imgsz
L4	5	32	640

따라서 작은 데이터셋을 활용해서 결과를 확인해보고 결정하는 것을 추천한다.

 

2. compress images

두 번째는 압축률을 이용하는 방법이다. 이미지 압축을 통해 이미지 하나하나의 크기를 줄일 수 있다. 프로젝트에서 이용한 데이터셋의 경우 500만 화소 이상의 고화질 이미지로 이미지 하나하나의 크기가 매우 컸다. 따라서 이미지 압축을 통해 크기를 줄였다. 이 방법 역시 작은 데이터셋을 이용해 테스트를 해보는 것이 좋다.

사진3, AI Hub-행동 분류 및 상호작용 인식용 한국형 비전 데이터

사진3을 보면 이미지를 압축했지만 차이가 없어 보인다. 오른쪽 사진이 원본이고 왼쪽이 압축한 사진이다. 하지만 이미지의 크기는 각각 309KB, 245KB이다.

사진4

위치가 살짝 다르지만 같은 배율만큼 확대했을 때다. 오른쪽이 원본, 왼쪽이 압축한 사진이다. 214% 확대했지만 큰 차이가 없다. 따라서 압축을 하는 것도 데이터의 크기를 줄이는 데 좋은 방법이라 생각한다. 물론 압축 강도에 대해서는 작은 데이터셋을 이용해서 확인을 해야 한다.

original images: 0.99006 0.80327
compress images: 0.9904	0.79492

GPU	epochs	batch size	imgsz
L4	5	32	640

압축하기 전과 후의 학습 결과는 mAP50, mAP50-95의 성능은 위와 같이 나왔다. 거의 차이가 없었다.