Yolo PyTorch custom dataset 학습시키기 [3]

JunSeong An·2022년 6월 29일
0

AI

목록 보기
3/4
post-thumbnail

이번 포스팅에서는 이전 포스팅에서 라벨링했던 정보를 yolov5s 모델에 넣어 학습시킬 예정입니다.


3. yolov5s로 학습 진행

yolov5s 같은 경우 추가적인 다운과정 없이 ultralytics github 페이지를 clone하시면 됩니다.

$ git clone git@github.com:ultralytics/yolov5.git

3-0 : dataset.yml

우선 현재는 다음과 같이 data 폴더와 yolov5폴더가 같이 디렉토리 안에 존재하고 있습니다. 우선 yolov5폴더에 dataset.yml을 추가해줍니다.

dataset.yml

# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: ../data  # dataset root dir
train: images  # train images (relative to 'path') 128 images
val: images  # val images (relative to 'path') 128 images

# Classes
nc: 17 # number of classes
names: [ 'dog', 
'person',
'cat',
'tv',
'car',
'meatballs',
'marinara sauce',
'tomato soup',
'chicken noodle soup',
'french onion soup',
'chicken breast',
'ribs',
'pulled pork',
'hamburger',
'cavity',
'scrollup',
'scrolldown'
]  # class names
  1. 우선 path 설정입니다. dataset.yml 파일은 yolov5 안에 존재하기 때문에 이전 폴더로 넘어가 data에 접근하기 때문에 ../data를 path로 넣어줍니다.
  2. 그리고 train data는 지정한 path의 images 폴더 안에 있기 때문에 imagse를 폴더이름으로 넣어줍니다
  3. val data는 학습한 데이터를 평가할 때 쓰이는데, 저희는 가볍게 학습시킬 것이기 때문에 따로 val data를 분리하지 않았습니다. train data와 동일한 폴더를 넣겠습니다.
  4. 이전 포스팅에서 classes.txt에 있던 객체의 개수를 nc의 값으로 넣어줍니다.

주의사항 : labeling을 할 때 진행했던 label의 이름과 names의 값을 넣어주세요!

이렇게 되면 dataset.yml파일은 준비되었습니다.

3-1 : 학습시키기

yolov5폴더 안의 train.py를 실행시키면 됩니다.

$ python train.py --img 320 --batch 16 --epochs 500 --data dataset.yml --weights yolov5s.pt --workers 2

저희가 정의한 dataset.yml파일을 꼭 넣어주세요!

이렇게 되면 다음과 같은 값이 뜨며 학습을 진행합니다.

그리고 학습된 .pt 파일은 yolov5 안에 runs/train/exp/weights 폴더 안에 best.pt와 last.pt파일을 사용하면 됩니다.


이제 학습까지 모두 끝났습니다. 다음 포스팅에서는 만들어진 .pt파일을 이용해 실제로 detection을 진행하고 어떻게 하면 활용할 수 있을지에 대해 포스팅 해보겠습니다.

profile
안녕하세요!

1개의 댓글

comment-user-thumbnail
2023년 7월 10일

전 개발환경이 Windows에 IDE는 Visual Studio Code를 쓰고 있는데요, $ python train.py --img 320 --batch 16 --epochs 500 --data dataset.yml --weights yolov5s.pt --workers 2 이 부분을 어떻게 처리해야할지 모르겠습니다. 터미널창에 $를 빼고 입력해봐도 뜨질 않아서,,,

답글 달기