본문 바로가기
딥러닝관련/Detection

YOLOv5 PyTorch 돌려보기 (작성 중)

by 머리올리자 2023. 1. 18.

https://github.com/ultralytics/yolov5

 

GitHub - ultralytics/yolov5: YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite

YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite. Contribute to ultralytics/yolov5 development by creating an account on GitHub.

github.com

위가 원작자 코드인가보다 STAR가 가장 많다.

 

https://github.com/ultralytics/yolov5/wiki/Docker-Quickstart

 

GitHub - ultralytics/yolov5: YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite

YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite. Contribute to ultralytics/yolov5 development by creating an account on GitHub.

github.com

위 링크로 들어가 docker image를 다운 받는다.

다운 완료 후 docker container를 실행 및 학습을 해본다.

sudo docker run --ipc=host -it --gpus all ultralytics/yolov5:latest
python train.py  # train a model

동작이 일단은 잘 된다.

 

터미널을 보아하니 coco 전체 이미지가 아닌 126 장의 이미지만 사용하는듯 하다.

 

전체 이미지를 다운 받아서 실행해보자.

 

그 전에 container 내 말고 local에 git clone하여 마운트해봐야겠다.

data-scripts 디렉토리에서 아래 명령어를 수행하여 coco2017 데이터셋을 다운 받는다.

 

bash get_coco.sh

 

기존의 docker를 종료 시키고 새롭게 docker container를 열고 로컬 파일을 마운트 한다.

sudo docker run -v $(pwd):/workspace --ipc=host -it --gpus all ultralytics/yolov5:latest

위와 같이 로컬 파일들을 마운트 후 아래의 명령어를 수행하면 학습이 잘 진행되는 것을 확인할 수 있다.

python train.py --data coco.yaml --weights yolov5s.pt --img 640