글 작성자: 만렙개발자

YOLOv5는 x(x-large), l(large), m(medium), s(small)의 4종류로 구성됩니다. 자세한 내용은 이전 글을 참조해주시고, 이 포스트에서는 학습 코드(스크립트)에 대해서 다루고자 합니다.

 

[분류 전체보기] - YOLO v5 공개! 개념, 이론, 논문 대신에 iOS...?

 

YOLO v5 공개! 개념, 이론, 논문 대신에 iOS...?

YOLO V5가 공개되었습니다! YOLO는 You Only Look Once라는 이름으로, one-stage object detection 딥러닝 기법으로 매우 빠른 속도의 추론 과정으로 2016년에 큰 충격을 안겨주었었습니다. 저는 증강/가상현실을.

lv99.tistory.com

먼저 YOLOv5를 official하게 공개한 ultralytics의 github repo로부터 clone을 하고, 필요한 requirements에 대해서 설치를 합니다. 그리고 torch, ipython display, google drive download 패키지들을 import 합니다.

 

!git clone https://github.com/ultralytics/yolov5 
!pip install -r yolov5/requirements.txt 
%cd yolov5

import torch
from IPython.display import Image, clear_output
from utils.google_utils import gdrive_download

clear_output()
print('Setup complete. Using torch %s %s' % (torch.__version__, torch.cuda.get_device_properties(0) if torch.cuda.is_available() else 'CPU'))

 

간단하게 학습을 테스트해보기 위해서 마스크를 쓴 얼굴 데이터셋을 이용해보고자 합니다. 아래의 링크에서 다운로드 가능합니다. roboflow는 데이터셋을 구성하는 서비스를 제공합니다. 이미지와 태깅 데이터를 입력하고, 학습/검증/시험  index 설정, 데이터 증강 등 다양한 기능을 제공합니다. 특히 유용한 기능 중 하나는 원하는 데이터셋 형식으로 export할 수 있다는 점입니다. 아쉽게도 한 데이터셋에 1000개 이상의 이미지로 구성되는 경우에는 유료입니다. 저희가 이번에 사용할 예제 데이터셋은 무료입니다 :) 회원가입은 하셔야합니다ㅠㅠ (준비된 다른 데이터셋이 있으면 됩니다만, 예제 형식을 확인하기 위해서는 이것을 사용하면 효과적입니다.)

 

https://public.roboflow.ai/object-detection/mask-wearing

 

Mask Wearing Object Detection Dataset

Download 149 free images labeled with bounding boxes for object detection.

public.roboflow.ai

 

위 링크로 이동하시면 위와 같은 화면이 나타납니다. Downloads의 416x416-balck-padding을 클릭하면, 아래와 같이 데이터셋에 대한 구성이 나타납니다. 149개의 이미지로 가볍게 테스트해보기 좋은 크기로 구성되어있습니다. 오른쪽 상단의 Download 버튼을 누르면, 아래와 같이 Format을 설정할 수 있도록 나타나는데, YOLO v5 PyTorch를 선택하면 됩니다!

 

다운로드한 데이터셋의 압출을 풀면  test, train, valid 폴더와 data.yaml 파일, 그리고 README 파일들이 주어집니다. 폴더 이름이 너무 길어서 mask-149로 변경하였습니다. 그리고 현재 프로젝트 폴더(train.py 파일이 있는 디렉토리)에 위치시켰습니다.

 

# Start tensorboard
%load_ext tensorboard
%tensorboard --logdir runs

학습 과정을 살펴볼 수 있도록 tensorboard를 실행시켜줍니다. 참조할 폴더는 runs입니다. 디폴트 포트는 6006 입니다.

 

!python train.py --img 149 --batch 16 --epochs 5 --data ./mask-149/data.yaml \
--cfg ./models/yolov5s.yaml --weights yolov5s.pt --name tutorial --nosave --cache

YOLOv5s 학습을 실행하는 스크립트는 위와같이 매우 단순합니다. weights는 사전학습된 yolov5s.pt를 사용하며, 없을 경우에는 다운로드가 자동적으로 됩니다. 데이터셋의 정보를 담고있는 ./mask-149/data.yaml 파일을 살펴봅시다. 

train: ../train/images
val: ../valid/images

nc: 2
names: ['mask', 'no-mask']

nc는 number of class로 클래스/카테고리의 수를 지정해주면 됩니다. 이 데이터셋에서는 mask와 no-mask 2개로 구성되어있기 때문에 nc: 2입니다. 

 

 

이제 위의 !python train.py를 통해 학습을 시작해봅시다. OOM 에러가 뜰 경우, 메모리가 부족한 것이니 파라미터가 적은 모델(s)을 사용하시거나, 배치 사이즈(--batch)를 줄이시면 됩니다.

 

              from  n    params  module                                  arguments                     
  0             -1  1      3520  models.common.Focus                     [3, 32, 3]                    
  1             -1  1     18560  models.common.Conv                      [32, 64, 3, 2]                
  2             -1  1     20672  models.common.Bottleneck                [64, 64]                      
  3             -1  1     73984  models.common.Conv                      [64, 128, 3, 2]               
  4             -1  1    161152  models.common.BottleneckCSP             [128, 128, 3]                 
  5             -1  1    295424  models.common.Conv                      [128, 256, 3, 2]              
  6             -1  1    641792  models.common.BottleneckCSP             [256, 256, 3]                 
  7             -1  1   1180672  models.common.Conv                      [256, 512, 3, 2]              
  8             -1  1    656896  models.common.SPP                       [512, 512, [5, 9, 13]]        
  9             -1  1   1905152  models.common.BottleneckCSP             [512, 512, 2]                 
 10             -1  1   1248768  models.common.BottleneckCSP             [512, 512, 1, False]          
 11             -1  1    130815  torch.nn.modules.conv.Conv2d            [512, 255, 1, 1]              
 12             -2  1         0  torch.nn.modules.upsampling.Upsample    [None, 2, 'nearest']          
 13        [-1, 6]  1         0  models.common.Concat                    [1]                           
 14             -1  1    197120  models.common.Conv                      [768, 256, 1, 1]              
 15             -1  1    313088  models.common.BottleneckCSP             [256, 256, 1, False]          
 16             -1  1     65535  torch.nn.modules.conv.Conv2d            [256, 255, 1, 1]              
 17             -2  1         0  torch.nn.modules.upsampling.Upsample    [None, 2, 'nearest']          
 18        [-1, 4]  1         0  models.common.Concat                    [1]                           
 19             -1  1     49408  models.common.Conv                      [384, 128, 1, 1]              
 20             -1  1     78720  models.common.BottleneckCSP             [128, 128, 1, False]          
 21             -1  1     32895  torch.nn.modules.conv.Conv2d            [128, 255, 1, 1]              
 22   [-1, 16, 11]  1         0  models.yolo.Detect                      [80, [[10, 13, 16, 30, 33, 23], [30, 61, 62, 45, 59, 119], [116, 90, 156, 198, 373, 326]]]
Model Summary: 165 layers, 7.07417e+06 parameters, 7.07417e+06 gradients

모델의 구성이 어떻게 되어있는지 Summary를 해주고,

Epoch gpu_mem GIoU obj cls total targets img_size 와 같은 데이터가 출력이 되면서 학습이 진행됩니다.

 

tensorboard로 확인할 수 있으며, 학습된 best와 last 모델은 weights 폴더에 저장됩니다.

 

학습데이터 어떤 식으로 증강되고, 정답이 어떻게 박스로 잡혔는지 train-bat1.jpg 파일을 통해 확인할 수 있습니다. 확인하는 가장 큰 이유 중에 하나는 박스가 제대로 안잡혔을 수 있기 때문입니다. roboflow가 다양한 format으로 export할 수 있다는 점이 장점인 이유가 여기에 있습니다. Format에 따라 바운딩 박스는 (x, y, w, h), x, y, x+w, y+h), (x_center, y_center, w, h) 등으로 표현됩니다. YOLO v5에서는 (x_center, y_center, w, h)의 형식을 따릅니다.

Image(filename='./train_batch1.jpg', width=900)  # view augmented training mosaics

다음번에는 Kaggle의 Global Wheat Detection에 YOLOv5 모델을 적용하여 학습하고 추론하는 것에 대해서 다뤄보도록 하겠습니다.