How to Train YOLOv5 On a Custom Dataset

원본 링크

How to Train YOLOv5 On a Custom Dataset

객체 탐지 모델의 YOLO 제품군은 Ultralytics의 YOLOv5 도입으로 더 강력해 졌다. 이글에서는 커스텀 사용을 위해 커스텀 객체를 인식하기 위한 YOLOv5를 훈련하는 방법을 살펴본다.

COCO 가중치로 라이브 비디오 YOLOv5 추론(inferencing) - 커스텀 YOLOv5 가중치를 사용하는지 알아보자.

이 저장소를 함께 통합한 것에 대해 Ultralytics에 감사한다. - 이는 데이터 관리 도구와 이 기술들이 컴퓨터 비젼 프로젝트를 배포하려는 개발자가 쉽게 접근할 수 있게 한다.

여기서는 public blood cell detection dataset - 이후 혈액데이터셋을 사용한다.

탐색기(detector)를 훈련하기 위해 다음 단계를 수행한다.

• YOLOv5 설치
• 커스텀 YOLOv5 객체 탐지 데이터 다운로드
• YOLOv5 모델 설정과 구조(architecture) 정의
• 커스텀 YOLOv5 탐색기 훈련
• YOLOv5 성능 평가
• YOLOv5 훈련 데이터 시각화
• 테스트 이미지로 YOLOv5 추론(inference)
• 이후 추론을 위한 YOLOv5 가중치 저장

YOLOv5: What's New?

불과 두달전(이 글은 2020년 6월에 작성) Google Brain의 EfficientDet 도입은 무척 흥분시켰고 breaking down EfficientDet를 작성했었다. 실시간 객체 탐지 영역에서 YOLO제품군의 명성을 가릴것이라 생각했지만 틀렸다.

3주동안 다크넷(Darknet) 프레임워크에 YOLOv4가 출시되고 breaking down the research in YOLOv4을 작성했다.

그리고 이 글을 작성하기 몇시간 전에 YOLOv5가 발표되었고 엄청나게 매끄럽다는걸 알았다. YOLOv5는 Ultralytics PyTorch framework로 작성되었다.(Pytorch...pytorch...ㅠ.ㅠ) 이것은 사용하기 매우 직관적이로 매우 빨리 추론한다. 사실 여러사람들이 더 가벼운 라이브러리로 더 빠른 추론을 하기 위해 YOLOv3와 YOLOv4 Darknet 가중치를 Ultralytics PyTorch 가중치로 변환했다.

YOLOv5가 YOLOv4보다 더 성능기준에 맞을까? 이것에 대해서는 조금 후에 더 알아볼 것이다. 하지만 YOLOv5 vs. YOLOv4에서 쉽게 예측할 수 있다.

YOLOv5 vs EfficientDet의 성능 (updated 6/23)
(source)

(? 그래프에서 YOLOv4는 안보이는데 비교를 하고 있네?? 어떤 그래프를 기준으로 한건지 모르겠다.) YOLOv4는 YOLOv5 저장소상 평가의 왼쪽에 잇다. 이는 이 글의 초기 실행 결과를 바탕으로 YOLOv5는 확실하게 사용하기 쉽고 커스텀 데이터상에서 매우 성능이 좋다는 것을 나타낸다.

진행하면서 YOLOv5 Colab Notebook에서 동시에 따라하는 것이 좋다.

YOLOv5 환경 설치

YOLOv5를 시작하기 위해 우선 YOLOv5 저장소를 복제(clone)하고 의존관계에 있는 것을 설치한다. 이것이 객체 탐지를 훈련과 추론 명령을 실행을 준비가 된 프로그래밍 환경을 설정한다.

!git clone https://github.com/ultralytics/yolov5  # clone repo
!pip install -U -r yolov5/requirements.txt  # install dependencies

%cd /content/yolov5

그러면, Google Colab에스 무료로 제공하는 훈련 환경을 볼 수 있다.

import torch
from IPython.display import Image  # for displaying images
from utils.google_utils import gdrive_download  # for downloading models/datasets

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

Google Colab에서 Tesla P100 GPU를 할당받을 가능성이 높다. 다음은 필자가 할달받은 리소스이다.

torch 1.5.0+cu101 _CudaDeviceProperties(name='Tesla P100-PCIE-16GB', major=6, minor=0, total_memory=16280MB, multi_processor_count=56)

GPU는 훈련시간을 가속시킨다. Colab은 또한 torch와 cuda가 미리 설치되어 있어서 좋다. 만약 로컬환경에서 시도한다면 YOLOv5를 설정하기 위한 추가적인 단계가 있을 것이다.

커스텀 YOLOv5 객체 탐지 데이터 다운로드

이 글에서는 Roboflow에서 커스텀 객체 탐지 데이터를 YOLOv5 포멧으로 다운로드한다. 이 글에서 public blood cell detection dataset으로 혈액 흐름에서 세포를 탐지하기 위해 YOLOv5를 훈련한다.

Labeling Tool 퀵노트

만약 레이블링되지 않은 이미지가 았다면, 우선 레이블링해야 한다. 무료 오픈소스 레이블링툴로 Roboflow Annotate 또는 getting started with LabelImg 또는 getting started with CVAT 어노테이션 도구의 가이드를 따르는 것이 좋다. 여기서는 처리하기 위해 50개까지의 이미지를 레이블링하였다. 이후 모델 성능을 개선을 위해 좀 더 많은 레이블이 필요할 것이다.

Note : 이 Roboflow Annotate라는 녀석이 별도로 다운로드 하는 것이 아니고, roboflow에 가입 후 데이터셋을 업로드 및 작업후 다시 다운로드 받는 형식이다.

Roboflow Annotate를 사용한 이미지 레이블링

Roboflow에서 계정을 생성하면 VOC XML, COCO JSON, TensorFlow Object Detection CSV 둥으로 데이터셋을 얻을 수 있다.

데이터셋 업로드 후 전처리와 보강 단계를 선택할 수 있다.

BCCD 샘플 데이터에 대한 설정

그리고 Generate와 Download를 누르면 YOLOv5 PyTorch format을 선책할 수 있다.

YOLOv5 PyTorch 선택

프롬프트가 나타나면 "Show Code Snippet"을 선택한다. 이것은 다운로드 curl script를 출력한다. 따라서 쉽게 데이터를 적절한 형태로 colab으로 이식할 수 있다.

curl -L "https://public.roboflow.ai/ds/YOUR-LINK-HERE" > roboflow.zip; unzip roboflow.zip; rm roboflow.zip

커스텀 데이터셋 다운로드

YOLOv5 포멧으로 커스텀 object 데이터셋 다운로드

Export는 YOLOv5 image, label 폴더 위치와 커스텀 클래스 정보를 나타내는 data.yaml아리 불리우는 YOLOv5.yaml 파일을 생성한다.

Define YOLOv5 Model Configuration and Architecture

다음으로 커스텀 객체 탐지를 위한 모델 설정 파일을 작성한다. 이 글에서는 가장 작고 가장 빠른 YOLOv5의 기초 모델을 사용한다. 아래를 포함하여 다른 YOLOv5 모델을 선택할 수 도 있다.

• YOLOv5s
• YOLOv5m
• YOLOv5l
• YOLOv5x

또한 비록 이 작업을 할 필요가 거의 없을지라도 이번 단계에서 네트워크의 구조를 수정할 수있다. 아래에 custom_yolov5s.yaml로 이름지은 YOLOv5 모델 설정이 있다.

nc: 3
depth_multiple: 0.33
width_multiple: 0.50

anchors:
  - [10,13, 16,30, 33,23] 
  - [30,61, 62,45, 59,119]
  - [116,90, 156,198, 373,326] 

backbone:
  [[-1, 1, Focus, [64, 3]],
   [-1, 1, Conv, [128, 3, 2]],
   [-1, 3, Bottleneck, [128]],
   [-1, 1, Conv, [256, 3, 2]],
   [-1, 9, BottleneckCSP, [256]],
   [-1, 1, Conv, [512, 3, 2]], 
   [-1, 9, BottleneckCSP, [512]],
   [-1, 1, Conv, [1024, 3, 2]],
   [-1, 1, SPP, [1024, [5, 9, 13]]],
   [-1, 6, BottleneckCSP, [1024]],
  ]

head:
  [[-1, 3, BottleneckCSP, [1024, False]],
   [-1, 1, nn.Conv2d, [na * (nc + 5), 1, 1, 0]],
   [-2, 1, nn.Upsample, [None, 2, "nearest"]],
   [[-1, 6], 1, Concat, [1]],
   [-1, 1, Conv, [512, 1, 1]],
   [-1, 3, BottleneckCSP, [512, False]],
   [-1, 1, nn.Conv2d, [na * (nc + 5), 1, 1, 0]],
   [-2, 1, nn.Upsample, [None, 2, "nearest"]],
   [[-1, 4], 1, Concat, [1]],
   [-1, 1, Conv, [256, 1, 1]],
   [-1, 3, BottleneckCSP, [256, False]],
   [-1, 1, nn.Conv2d, [na * (nc + 5), 1, 1, 0]],

   [[], 1, Detect, [nc, anchors]],
  ]

Training Custom YOLOv5 Detector

data.yaml과 custom_yolov5.yaml로 훈련할 준 준비가 되었다.
훈련을 시작하기 위해 아래 옵션으로 훈련을 할 수 있다.

• img : 입력 이미지 크기 정의
• batch : 배치 사이즈 결정
• epochs : 훈련 에포크의 수 정의(Note : 주로 3000이상)
• data : yaml 파일의 경로 설정
• cfg : 모델 설정 지정
• weights : 가중치 custom 경로 지정(Note : Ultralytics Google Drive folder에서 다운로드 가능)
• name : 결과명
• nosave : 최종 checkpoint만 저장
• cache : 더 빠른 훈련을 위한 캐싱 이미지

훈련을 시행한다.

커스텀 YOLOv5 탐지기 훈련. 빠르게 훈련한다.

커스텀 YOLOv5 탐지기 성능 평가

훈련이 완료되었다. 검즘 지표로 훈련 절차가 얼마자 잘 수행되었는지를 평가할 수 있다. 훈련 스크립트는 runs에 텐스플로 로그를 남긴다.

커스텀 데이터셋에서 텐서보드 결과 시각화

그리고 어떠한 이유에서 텐서보드를 시각화 할 수 없다면 utils.plot_results로 그래프를 그리고 result.png로 저장할 수도 있다.

여기서는 좀 일찍 훈련이 멈췄다. 검증 mAP가 가장 높은 훈련된 가중치를 가져온다.

YOLOv5 훈련 데이터 시각화

훈련동안, YOLOv5 훈련 파이프라인은 증강(augmentation)된 훈련 데이터의 배치(batch)를 생성한다. 훈련 데이터 실측값은 물론 증간된 훈련데이터를 시각화 할 수 있다.

훈련 데이터 실측치(ground truth)

YOLOv5로 증강된 훈련 데이터

Run YOLOv5 Inference on Test Images

이제 모델을 훈련하고 테스트이미지에서 추론을 수행했다. 훈련이 완료된 후 메델 가중치는 weight/ 경로에 저장된다. 추론을 위해 모델 신뢰도(confidence, 높은 신뢰도는 요구는 낮은 예측을 만든다.)를 나타내는 conf와 추론 source를 이들 가중치와 함께 호출한다. source는 이미지 디렉토리, 개별 이미지, 비디오 파일 그리고 웹캠 또한 가능하다. 소스의 경우 test/*jpg를 test_infer/로 이동시켰다.

!python detect.py --weights weights/last_yolov5s_custom.pt --img 416 --conf 0.4 --source ../test_infer

추론시간은 아주 빠르다. Tesla P100에서 YOLOv5는 142FPS에 달한다.

142FPS(.007s/image)에 도달한 YOLOv5 추론

마지막으로 테스트이미지에서 탐지기 추론 결과를 시각화해본다.

테스트 이미지에서의 YOLOv5 추론

Export Saved YOLOv5 Weights for Future Inference

이제 커스텀 YOLOv5 객체 탐지기를 검증했다. 이 모델을 실시간 컴퓨터 비젼 작업을 위해 Colab에서 가중치를 가져오고 싶을 것이다. 이를 위해서 Google Dirve 모듈을 임포트하여 내보내면 된다.

from google.colab import drive
drive.mount('/content/gdrive')

%cp /content/yolov5/weights/last_yolov5s_custom.pt /content/gdrive/My\ Drive