[YOLOv7] custom training

colab 환경

0. 설치

import sys
import torch
print(f"Python version: {sys.version}, {sys.version_info} ")
print(f"Pytorch version: {torch.__version__} ")

!nvidia-smi

!# Download YOLOv7 code
!git clone <https://github.com/WongKinYiu/yolov7>
%cd yolov7
!ls

!pip install -r requirements.txt

1. 데이터 분리

• train, val 파일 생성

 

• 각각의 폴더 내부에 images, labels 폴더 생성

 

• 각각 이미지와 라벨링 된 텍스트 파일을 넣는다. 이때 이미지와 텍스트 파일의 이름이 같아야 한다.

 

2. 데이터 옮기기

나눈 데이터를 압축하여 yolov7/data 폴더로 옮긴다.

아래 코드를 통해 데이터 압축을 해제한다.

!cd data
!unzip -uq "/content/yolov7/data/yolo_data.zip" -d "/content/yolov7/data/"

최종적으로 yolov7/data 폴더 내부에는 val, train 폴더가 존재해야한다.

 

3. 설정 파일

• custom_data.yaml 파일

train: ./data/train
val: ./data/val

# number of classes
nc: 80

# class names
names: [ 'label', 'cap', 'ring' ]

이 파일을 yolov7/data 폴더에 넣는다.

• yolov7-custom.yaml 파일

# parameters
nc: 3  # number of classes
depth_multiple: 1.0  # model depth multiple
width_multiple: 1.0  # layer channel multiple

yolov7/cfg/traing 폴더의 yolov7-tiny 파일을 변형한 파일로 단순히 nc를 클래스 개수로 바꿨다.

이 파일을 yolov7/cfg/training 폴더 내부에 넣어준다.

 

4. 가중치 파일 다운로드

yolov7의 github releases에서 가중치 파일을 다운 받는다.

나는 yolov7-tiny로 진행하기 때문에 yolov7-tiny.pt 파일 다운로드했다.

다운로드 된 파일을 yolov7 폴더에 넣는다.

 

5. 학습 시작

!python train.py --workers  1 --device 0 --batch-size 8 --epochs 100 --img 640 640 --data data/custom_data.yaml --hyp data/hyp.scratch.custom.yaml --cfg cfg/training/yolov7-custom.yaml --name yolov7-custom --weights yolov7-tiny.pt

 

6. 가중치 파일 확인하기

yolov7/runs/train/yolov7-custom/weights/best.pt 파일이 생성된 것을 알 수 있다.

 

7. 추론

!# Detection
!python detect.py --weights /content/yolov7/runs/train/yolov7-custom2/weights/best.pt --conf 0.25 --img-size 640 --source  0084.jpg