YOLO Docker 환경 구축 및 학습하기

YOLO는 객체 탐지(Object Detection) 분야에서 널리 사용되는 모델 중 하나이다.

본 글에서는 Ultralytics에서 제공하는 공식 Docker 이미지를 기반으로 YOLO 학습 환경을 구축하고, train 하는 방법을 소개하겠다.

1. Docker Hub에서 YOLO 이미지 다운로드

Ultralytics는 공식적으로 Docker Hub에 YOLO 이미지를 제공한다.

해당 이미지를 통해 손쉽게 학습 환경을 구성할 수 있다.

• Docker Hub 링크:
  • https://hub.docker.com/u/ultralytics

간단하게 예시로 ultralytics/yolov5를 pull 해보겠다.

YOLOv5 이상 v8~v12 이상의 최신 버전을 쓰고 싶다면, 도커 컨테이너에서 그냥 pip install ultralytics를 하면 된다.

https://hub.docker.com/r/ultralytics/yolov5

 

• 이미지 다운로드 명령어:

docker pull ultralytics/yolov5

다운로드가 완료되면 로컬 환경에 ultralytics/yolov5 이미지가 추가된다.

---

2. Docker 컨테이너 실행

이미지를 기반으로 컨테이너를 실행할 수 있다.

docker run --ipc=host -v {local_dir}:{docker_dir} --gpus all -it ultralytics/yolov5:latest

예시:

• -ipc=host: PyTorch의 shared memory 이슈를 방지하기 위한 설정
• v: 로컬 디렉토리와 컨테이너 디렉토리를 연결하여 데이터 연동이 가능
• -gpus all: 모든 GPU 자원을 컨테이너에 할당
• it: 대화형(interactive) 모드로 bash 쉘에 진입

도커 컨테이너를 켜기만 하면 환경 세팅이 끝난다!

이제 컨테이너 내에서 train.py를 실행하여 커스텀 학습을 진행할 수 있다.

python3 train.py --data data/custom_dataset.yaml --epochs 10