Computer VIsion/응용(실전편)

PC에서 그래픽카드로 Yolov5 커스텀 트레이닝 해보기!

JackSmith 2024. 8. 31.

이번 포스팅부터 Opencv부터 yolo등을 활용하여 다양한 컴퓨터비전을 실습해보도록 하겠습니다. 프로젝트를 진행할때 바로 적용해볼 수 있을 만한 그러한 영양가있는 실습 위주로 포스팅을 진행해보겠습니다.

 

PC에서 그래픽카드로 Yolov5 커스텀 트레이닝 해보기 시작!

 

1. 환경(Environment)

 

저는 미니콘다 3.9버전을 다운받아, 미니콘다 가상환경을 기반으로 작업환경을 맞춰주었습니다.아래 링크를 타고 들어가줍니다.

 

Latest Miniconda installer links by Python version — Anaconda documentation

Latest Miniconda installer links by Python version This list of Miniconda installers is for all supported versions of Python, separated by operating system. For an archive of Miniconda versions, see https://repo.anaconda.com/miniconda/. Windows installers

docs.anaconda.com

그리고 페이지 하단에 있는 3.9버전을 설치해줍니다.

 

또한, 제가 갖고 있는 컴퓨터 사양을 확인해보면,

 

다음과 같이 그래픽카드 2070 Super를 사용하고 있기 때문에, 이를 바탕으로 nvidia 드라이버와 cuda를 설치해주도록 하겠습니다.

먼저 nvidia 드라이버를 설치해주도록 하겠습니다.

 

 

최신 정식 NVIDIA 드라이버 다운로드

최신 정식 NVIDIA 드라이버를 다운로드하여 PC 게임 경험을 향상시키고 앱을 더 빠르게 실행하세요.

www.nvidia.com

 

 

 

이렇게 nvidia 드라이버를 설치해주셨다면, 아래 명령어를 통해 정보를 입력받아 올수있는지 확인해보겠습니다.

nvidia-smi

 

 

이번에는 쿠다를 설치해주도록 하겠습니다. 아래 링크에 들어가줍니다.

 

CUDA Toolkit 12.1 Downloads

Get the latest feature updates to NVIDIA's proprietary compute stack.

developer.nvidia.com

 

 

이렇게 cuda를 설치해주셨다면, 아래 명령어를 통해 정보를 입력받아 올 수 있는지 확인해보겠습니다.

nvcc --version

이렇게 까지 했다면, GPU를 사용할수있도록 물려주었다고 볼 수 있습니다.

(뒤에서 한번더 파이썬환경에서 GPU자원을 끌어쓸 수 있는지 최종확인하는 방법을 알려드리겠습니다)

 

 

2. 코드작업하기

일단 저는 문서 디렉토리 아래에 아래와 같이 구조를 잡아주었습니다.

documents>playground>python-workspace>training

 

 

2.1 미니콘다 가상환경 생성하고 실행하기

이후 미니콘다 프롬프트를 켜서 해당 경로로 들어가주고 아래 명령어를 통해 가상환경을 생성해줍니다.(저는 가상환경 이름을 'elecbro'으로 잡아주었습니다)

conda create -n elecbro python=3.9.12

그리고 이 가상 인터프리터를 사용하기 위해 해당 가상환경으로 접속해줍니다.

conda activate elecbro

 

 

 

2.2 torch라이브러리를 가져와 GPU끌어올 수 있게 세팅해주기

이제 저희가 사용하려는 yolov5 트레이닝을 시켜주기 위해서는 파이토치 라이브러리가 필요합니다.

pip install torch==1.8.1+cu101 torchvision==0.9.1+cu101 torchaudio==0.8.1 -f https://download.pytorch.org/whl/torch_stable.html

 

이후 미니콘다 가상환경 내에서 python 스크립트창을 열어주고 아래와 같이 코드를 쳐줍니다.

>>> import torch
>>> print(torch.cuda.is_available())

이후 GPU가 제대로 물려있다면 위와 같이 "True"를 출력할 것입니다.

2.3 훈련 데이터셋 가져오기

아래 링크로 들어가면 Synthetic Fruit Dataset이 있을 겁니다.

 

Synthetic Fruit Object Detection Dataset

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

public.roboflow.com

 

여기서 아래와 같이 다운로드 버튼을 눌러주고,

아래와 같이 Format을 Yolov5 torch으로 잡아주도록 하겠습니다. 그리고 Continue 클릭!!

 

그러면 이제 위와 같은 창이 뜰텐데,  중간에 있는 curl명령문구를 복사해줍니다.

 

 

그리고 이전의 미니콘다 프롬프트 참으로 가서 해당 명령어를 붙여넣어주면 압축파일을 다운로드 받을 수 있습니다.

 

다운을 다 받았다면 아래와 같이 압축파일이 있을 것입니다.

압축해제 후, 저 roboflow안에 있는 데이터셋들을 모두 training 디렉토리 바로 아래로 이동시켜 줍니다.

before>

 

after>

(이제 roboflow폴더랑 압축파일은 필요없으니 지워주겠습니다ㅎㅎ)

 

2.4 코드 레포지토리 가져오기

아래 명령어를 통해 yolov5 레포지토리를 가져와줍니다.

git clone https://github.com/ultralytics/yolov5.git

 

이후 위 레포지토리를 실질적으로 이용하기 위해 아래에 물려주어야 할 의존성 패키지를 설치해주어야 합니다.

cd yolov5
pip install -r requirements.txt

 

 

3. Traing 시작하기

의존성 패키지들을 다 설치해주었다면, 다시 미니콘다 프롬프트(터미널)에서 training디렉토리로 올라와줍니다.

그리고 아래 명령어를 통해 훈련을 시켜줍니다.

python ./yolov5/train.py --img 416 --batch 16 --epochs 50 --data ./data.yaml --cfg ./yolov5/models/yolov5s.yaml --weights yolov5s.pt --name fruit_yolov5s_results

 

이후 제 그래픽카드( GTX 2070 Super )로 학습을 진행했을때는 15분정도가 소요되었습니다.

완성된 모델이 제대로 존재하는지 확인해줍니다.

경로는 yolov/runs/train/fruit_yolov5_results에 있습니다.

fruit_yolov5s_results 디렉토리 안에 들어가면 weight 디렉토리가 있는데 이 안에 best.pt라는 완성된 모델이 있습니다.

이제 이 모델을 활용해서 실제로 잘 동작하는지 확인해보도록 하겠습니다.

 

4. 모델 검증 및 테스트하기(Validation & Test)

일단 training디렉토리 바로 아래에 테스트용 이미지를 넣어줍니다. 테스트할 이미지는 아래의 remon.jpg입니다.

(해당 이미지를 저장하여 사용해보세요!)

 

아래 명령어를 통해 detect를 진행해보도록 하겠습니다.

python ./yolov5/detect.py --weights ./yolov5/runs/train/fruit_yolov5s_results/weights/best.pt --conf 0.5 --source ./remon.jpeg

위 로그를 보면 정상적으로 탐지를 한것 같습니다. "Results saved to yolov5\runs\detect\exp2"를 보아 이곳으로 들어가면 그 output을 확인해볼 수 있을 것 같습니다.

위와 같이 정상적으로 객체를 탐지해낸 것을 알 수 있습니다.

 

긴글 읽어주셔서 감사합니다!:):)

 

 

출처:https://www.youtube.com/watch?v=U-wkRQ8U3GE&t=303s

 

댓글