[YOLO] Install, Train, and Inference on Ubuntu

YOLO는 object detection에서 가장 많이 쓰이는 실시간 detector 계열이다. 이미지나 영상에서 객체의 class와 bounding box를 한 번에 예측한다.

이 글은 Ubuntu 기준으로 YOLO 설치, 학습, 추론을 정리한다. Ultralytics YOLO를 쓰면 설치, 학습, 검증, 추론, export를 CLI와 Python API로 모두 처리할 수 있다. 공식 문서의 현재 예시는 yolo26n.pt를 사용하지만, yolo11n.pt, yolov8n.pt 같은 기존 모델도 같은 흐름으로 사용할 수 있다.

YOLO 실습의 전체 흐름은 설치 → 데이터셋 구성 → 학습 → 검증 → 추론 → export 순서로 보면 된다.

Ubuntu 준비

먼저 Ubuntu package를 업데이트한다.

sudo apt update
sudo apt upgrade -y

Python 가상환경과 이미지/동영상 처리를 위한 기본 패키지를 설치한다.

sudo apt install -y \
  python3 \
  python3-pip \
  python3-venv \
  git \
  curl \
  wget \
  libgl1 \
  libglib2.0-0

libgl1, libglib2.0-0은 OpenCV가 이미지를 읽거나 저장할 때 Ubuntu에서 자주 필요한 system library다. 서버 환경에서 ImportError: libGL.so.1 같은 오류가 나면 보통 이 계열 패키지가 빠진 것이다.

NVIDIA GPU가 있다면 driver가 잡히는지 먼저 확인한다.

nvidia-smi

정상이라면 GPU 이름, driver version, CUDA version 정보가 출력된다. nvidia-smi가 없거나 실패하면 YOLO 설치보다 NVIDIA driver 설치를 먼저 해결해야 한다.

Ubuntu에서 YOLO 학습을 제대로 하려면 Python package보다 먼저 NVIDIA driver와 nvidia-smi 상태를 확인해야 한다.

YOLO 설치

Python 가상환경을 먼저 만든다.

mkdir -p ~/Projects/yolo-practice
cd ~/Projects/yolo-practice

python3 -m venv .venv
source .venv/bin/activate
python -m pip install -U pip

일반 desktop Ubuntu라면 Ultralytics package를 설치한다.

pip install -U ultralytics

서버나 Docker처럼 display가 없는 Ubuntu 환경에서는 OpenCV GUI 의존성을 피하기 위해 headless package를 쓸 수 있다.

pip install ultralytics-opencv-headless

둘 중 하나만 선택하면 된다. 일반 노트북/데스크탑이면 ultralytics, headless server면 ultralytics-opencv-headless가 무난하다.

설치 확인:

yolo version

또는 Python에서 확인한다.

python - <<'PY'
from ultralytics import YOLO
print("ultralytics import ok")
PY

GPU를 쓸 계획이라면 PyTorch CUDA 설치 상태도 확인한다.

python - <<'PY'
import torch
print(torch.__version__)
print("cuda:", torch.cuda.is_available())
if torch.cuda.is_available():
    print(torch.cuda.get_device_name(0))
PY

CUDA가 False이면 학습은 CPU로도 가능하지만 매우 느리다. 실제 학습은 NVIDIA GPU 환경에서 하는 것이 좋다.

Ubuntu에서 CUDA가 잡히지 않는 경우에는 다음 순서로 확인한다.

nvidia-smi
which python
python -m pip show torch
python - <<'PY'
import torch
print(torch.version.cuda)
print(torch.cuda.is_available())
PY

nvidia-smi는 되는데 torch.cuda.is_available()가 False이면 PyTorch가 CPU build로 설치되었을 가능성이 크다. 이때는 PyTorch 공식 설치 명령으로 CUDA 지원 wheel을 다시 설치한다.

CLI 기본 구조

Ultralytics CLI는 다음 형식을 사용한다.

yolo TASK MODE ARGS

주의할 점은 argument 형식이다.

yolo detect predict model=yolo26n.pt source=image.jpg imgsz=640 conf=0.25

Ultralytics CLI는 --model 같은 argparse 스타일이 아니라 key=value 형식을 쓴다.

YOLO CLI argument는 key=value 형식으로 써야 하며, --model, --imgsz 같은 prefix를 붙이지 않는다.

빠른 추론 테스트

설치가 끝났으면 pretrained model로 바로 추론할 수 있다.

yolo detect predict model=yolo26n.pt source='https://ultralytics.com/images/bus.jpg'

로컬 이미지:

yolo detect predict model=yolo26n.pt source=images/test.jpg

폴더 전체:

yolo detect predict model=yolo26n.pt source=images/

동영상:

yolo detect predict model=yolo26n.pt source=videos/input.mp4

웹캠:

yolo detect predict model=yolo26n.pt source=0

결과는 기본적으로 runs/detect/predict* 아래에 저장된다. confidence threshold를 조절하려면 conf를 쓴다.

yolo detect predict model=yolo26n.pt source=images/ conf=0.4 imgsz=640

데이터셋 구조

Object detection dataset은 보통 이미지와 label text file로 구성한다. YOLO label은 이미지마다 같은 이름의 .txt 파일을 둔다.

datasets/
  my_dataset/
    images/
      train/
        img001.jpg
        img002.jpg
      val/
        img101.jpg
        img102.jpg
    labels/
      train/
        img001.txt
        img002.txt
      val/
        img101.txt
        img102.txt
    data.yaml

label 파일은 한 줄에 객체 하나를 기록한다.

class_id x_center y_center width height

예시:

0 0.512 0.438 0.231 0.184
1 0.221 0.610 0.104 0.160

여기서 좌표는 pixel 값이 아니라 0부터 1 사이로 normalize된 값이다.

YOLO detection label은 class x_center y_center width height 순서이며, box 좌표는 pixel이 아니라 0~1 범위의 normalized xywh 값이다.

객체가 없는 이미지에는 label file을 만들지 않아도 된다. 다만 이미지와 label 이름이 맞지 않으면 학습 중 missing label 문제가 생길 수 있다.

data.yaml 작성

data.yaml은 dataset root, train/val 경로, class 이름을 정의한다.

path: datasets/my_dataset
train: images/train
val: images/val
test:

names:
  0: person
  1: helmet
  2: drone

path를 기준으로 train, val 경로가 해석된다. 위 예시는 다음 경로를 의미한다.

datasets/my_dataset/images/train
datasets/my_dataset/images/val

class id는 반드시 0부터 시작해야 한다. label file에 2가 있는데 names에 2가 없으면 학습이 실패하거나 class mapping이 꼬인다.

학습

Ubuntu terminal에서 가상환경을 활성화한 뒤 학습한다.

cd ~/Projects/yolo-practice
source .venv/bin/activate

가장 간단한 학습 명령은 다음과 같다.

yolo detect train data=datasets/my_dataset/data.yaml model=yolo26n.pt epochs=100 imgsz=640

GPU를 명시하려면 device를 쓴다.

yolo detect train data=datasets/my_dataset/data.yaml model=yolo26n.pt epochs=100 imgsz=640 device=0

여러 GPU 중 1번 GPU를 쓰려면 다음처럼 지정한다.

yolo detect train data=datasets/my_dataset/data.yaml model=yolo26n.pt epochs=100 imgsz=640 device=1

2개 GPU를 쓰는 예시는 다음과 같다.

yolo detect train data=datasets/my_dataset/data.yaml model=yolo26n.pt epochs=100 imgsz=640 device=0,1

CPU로 강제 실행:

yolo detect train data=datasets/my_dataset/data.yaml model=yolo26n.pt epochs=10 imgsz=640 device=cpu

batch size를 지정할 수도 있다.

yolo detect train data=datasets/my_dataset/data.yaml model=yolo26n.pt epochs=100 imgsz=640 batch=16

GPU memory가 부족하면 batch size를 줄인다.

yolo detect train data=datasets/my_dataset/data.yaml model=yolo26n.pt epochs=100 imgsz=640 batch=8

그래도 OOM이 나면 imgsz를 줄이거나 더 작은 모델을 사용한다.

yolo detect train data=datasets/my_dataset/data.yaml model=yolo26n.pt epochs=100 imgsz=512 batch=8

학습 결과는 기본적으로 다음 위치에 저장된다.

runs/detect/train/

중요한 파일은 다음 두 개다.

runs/detect/train/weights/best.pt
runs/detect/train/weights/last.pt

일반적으로 추론과 배포에는 best.pt를 먼저 사용한다.

학습 재개

학습이 중간에 끊겼다면 resume=True를 사용한다.

yolo detect train model=runs/detect/train/weights/last.pt resume=True

새로운 설정으로 이어서 fine-tuning하고 싶다면 resume=True보다 model=...에 checkpoint를 넣고 새 학습으로 시작하는 편이 명확하다.

yolo detect train data=datasets/my_dataset/data.yaml model=runs/detect/train/weights/best.pt epochs=50 imgsz=640

검증

학습된 모델을 validation set에서 평가한다.

yolo detect val model=runs/detect/train/weights/best.pt data=datasets/my_dataset/data.yaml imgsz=640

출력에서 주로 보는 값은 precision, recall, mAP이다.

데이터가 적거나 label 품질이 나쁘면 training loss는 내려가도 validation mAP가 낮을 수 있다. 이 경우 모델보다 dataset split, annotation quality, class imbalance를 먼저 확인하는 것이 좋다.

추론

Ubuntu terminal에서 가상환경을 활성화한다.

cd ~/Projects/yolo-practice
source .venv/bin/activate

학습한 best.pt로 이미지 추론:

yolo detect predict model=runs/detect/train/weights/best.pt source=images/test.jpg conf=0.25

폴더 추론:

yolo detect predict model=runs/detect/train/weights/best.pt source=images/test/ conf=0.25

동영상 추론:

yolo detect predict model=runs/detect/train/weights/best.pt source=videos/input.mp4 conf=0.25

결과 저장 위치를 지정하려면 project, name을 사용한다.

yolo detect predict \
  model=runs/detect/train/weights/best.pt \
  source=images/test/ \
  conf=0.25 \
  project=runs/inference \
  name=my_dataset_test

결과는 다음처럼 저장된다.

runs/inference/my_dataset_test/

Python API로 학습과 추론

CLI는 빠른 실험에 좋고, Python API는 서비스나 pipeline에 넣기 좋다.

학습 코드

train_yolo.py:

from ultralytics import YOLO


def main() -> None:
    model = YOLO("yolo26n.pt")

    model.train(
        data="datasets/my_dataset/data.yaml",
        epochs=100,
        imgsz=640,
        batch=16,
        device=0,
        project="runs/detect",
        name="my_dataset_yolo26n",
    )


if __name__ == "__main__":
    main()

실행:

python train_yolo.py

Windows에서는 multiprocessing 문제를 피하기 위해 if __name__ == "__main__": guard를 넣는 것이 중요하다. Linux에서도 이 형태로 쓰는 습관을 들이는 것이 좋다.

추론 코드

predict_yolo.py:

from pathlib import Path

from ultralytics import YOLO


def main() -> None:
    model = YOLO("runs/detect/my_dataset_yolo26n/weights/best.pt")
    results = model.predict(
        source="images/test",
        conf=0.25,
        imgsz=640,
        save=True,
        project="runs/inference",
        name="python_predict",
    )

    for result in results:
        print("image:", Path(result.path).name)
        print("boxes:", len(result.boxes))


if __name__ == "__main__":
    main()

동영상이나 stream처럼 입력이 길면 stream=True를 사용해 결과를 generator로 처리하는 것이 좋다.

from ultralytics import YOLO


model = YOLO("runs/detect/my_dataset_yolo26n/weights/best.pt")

for result in model.predict(source="videos/input.mp4", stream=True, conf=0.25):
    print(result.path, len(result.boxes))

긴 동영상이나 실시간 stream 추론에서는 stream=True를 사용해 결과를 한 프레임씩 처리하는 것이 메모리 관리에 유리하다.

Export

배포를 위해 ONNX로 export할 수 있다.

yolo export model=runs/detect/train/weights/best.pt format=onnx imgsz=640

Python API:

from ultralytics import YOLO

model = YOLO("runs/detect/train/weights/best.pt")
model.export(format="onnx", imgsz=640)

TensorRT를 쓸 환경이라면 format=engine도 사용할 수 있다. 다만 TensorRT export는 CUDA, TensorRT version, GPU driver에 민감하므로 서버 환경을 먼저 고정하는 것이 좋다.

Ubuntu GPU 서버에서 TensorRT export를 할 때는 먼저 NVIDIA driver와 TensorRT 설치 상태를 확인한다.

nvidia-smi
python - <<'PY'
import torch
print(torch.cuda.is_available())
PY

추천 프로젝트 구조

실험을 반복할 때는 데이터, 코드, 결과를 분리한다.

yolo-practice/
  datasets/
    my_dataset/
      images/
        train/
        val/
      labels/
        train/
        val/
      data.yaml
  scripts/
    train_yolo.py
    predict_yolo.py
    export_yolo.py
  runs/
    detect/
    inference/
  README.md

runs/는 실험 결과가 계속 쌓이는 폴더다. Git으로 관리할 때는 보통 제외한다.

runs/
*.pt
*.onnx

모델 weight는 용량이 크므로 Git 대신 object storage, release asset, experiment tracking tool로 관리하는 편이 낫다.

자주 나는 문제

ModuleNotFoundError: No module named 'ultralytics'

가상환경이 활성화되지 않았거나 다른 Python에 설치했을 가능성이 크다.

which python
which pip
pip show ultralytics

CUDA가 안 잡힘

먼저 Ubuntu에서 NVIDIA driver가 정상인지 확인한다.

nvidia-smi

PyTorch가 CUDA build로 설치되어 있는지 확인한다.

python - <<'PY'
import torch
print(torch.__version__)
print(torch.cuda.is_available())
PY

False이면 NVIDIA driver, CUDA 지원 PyTorch wheel, 가상환경을 확인한다.

ImportError: libGL.so.1

Ubuntu server나 Docker image에서 OpenCV가 GUI 관련 system library를 찾지 못할 때 자주 난다.

sudo apt update
sudo apt install -y libgl1 libglib2.0-0

GUI가 필요 없는 서버라면 headless package를 쓰는 것도 방법이다.

pip uninstall -y ultralytics
pip install ultralytics-opencv-headless

Permission denied

dataset이나 runs/ 폴더를 root 권한으로 만들어서 현재 사용자가 쓰지 못하는 경우가 있다.

ls -ld datasets runs

필요하면 소유권을 현재 사용자로 바꾼다.

sudo chown -R "$USER:$USER" datasets runs

학습을 sudo yolo ...로 실행하는 것은 피하는 것이 좋다. 가상환경과 cache 경로가 꼬일 수 있다.

GPU memory 부족

CUDA out of memory가 나면 batch size를 줄인다.

yolo detect train data=datasets/my_dataset/data.yaml model=yolo26n.pt epochs=100 imgsz=640 batch=4 device=0

그래도 부족하면 image size를 줄인다.

yolo detect train data=datasets/my_dataset/data.yaml model=yolo26n.pt epochs=100 imgsz=512 batch=4 device=0

label 좌표가 이상함

YOLO label은 normalized xywh다. pixel 좌표 (xmin, ymin, xmax, ymax)를 그대로 넣으면 학습이 망가진다.

class id가 맞지 않음

data.yaml의 names와 label file의 class id가 일치해야 한다. class id는 0부터 시작한다.

학습 결과가 너무 낮음

모델 크기를 키우기 전에 다음을 먼저 확인한다.

1. train/val split이 제대로 되었는가
2. label 누락이나 잘못된 box가 없는가
3. class imbalance가 심하지 않은가
4. validation set이 실제 테스트 환경과 비슷한가
5. image size가 작은 객체를 보기에 충분한가

정리

YOLO 실습은 다음 명령만 이해해도 기본 흐름을 잡을 수 있다.

sudo apt update
sudo apt install -y python3 python3-pip python3-venv git libgl1 libglib2.0-0

python3 -m venv .venv
source .venv/bin/activate
pip install -U ultralytics

yolo detect train data=datasets/my_dataset/data.yaml model=yolo26n.pt epochs=100 imgsz=640

yolo detect val model=runs/detect/train/weights/best.pt data=datasets/my_dataset/data.yaml

yolo detect predict model=runs/detect/train/weights/best.pt source=images/test/

yolo export model=runs/detect/train/weights/best.pt format=onnx

처음에는 CLI로 빠르게 실험하고, 반복 실험이나 서비스 통합이 필요해지면 Python API로 옮기는 방식이 좋다.

References

• Ultralytics Quickstart
• Ultralytics CLI Guide
• Ultralytics Train Mode
• Ultralytics Predict Mode
• Ultralytics Object Detection Dataset Format