← 목록으로
YOLO26 튜토리얼: RTX 5060 8GB로 30분 만에 안전장비 탐지 모델 학습하기
gcube · hands-on

YOLO26 튜토리얼: RTX 5060 8GB로 30분 만에 안전장비 탐지 모델 학습하기

on this page

들어가며

GPU 클라우드를 떠올리면 보통 거대 언어 모델을 돌리는 고가의 GPU를 생각합니다. 하지만 현장에서 바로 쓰이는 AI 중에는 8GB GPU 한 장으로 충분한 것도 많습니다. 대표적인 예가 사진·영상 속 사물을 찾아내는 **객체 탐지(Object Detection)**입니다.

이 글에서는 gcube의 RTX 5060 8GB 컨테이너에서 Ultralytics의 최신 모델 YOLO26을 학습해, 공사 현장 사진에서 안전모·조끼·장갑·안전화·보안경 착용 여부를 찾는 모델을 만듭니다.

이 글을 끝까지 따라 하면

  • 사진을 넣으면 안전장비를 박스로 표시하는 나만의 모델(best.pt)을 얻습니다.
  • 학습 그래프와 정확도 지표(mAP)를 읽는 법을 익힙니다.
  • 다른 서버나 엣지 기기에서 쓸 수 있는 ONNX 파일까지 만듭니다.
항목내용
난이도초급 (Python 기초만 있으면 됨)
GPUgcube RTX 5060 8GB × 1
모델Ultralytics YOLO26n (선택: YOLO26s)
데이터Construction-PPE (1,416장, 11개 클래스, 자동 다운로드)
소요 시간약 30분 (YOLO26n 학습 7.5분 포함)

1. 미리 알아두기

YOLO26

  • Ultralytics가 2026년 1월 정식 출시한 최신 YOLO 모델입니다.
  • 기존 YOLO는 겹치는 박스를 지우는 후처리(NMS)가 따로 필요했지만, YOLO26은 이를 모델 안에서 끝내는 NMS-free 구조라 배포가 더 간단합니다.
  • 크기는 n · s · m · l · x 5종이며, 여기서는 8GB에 여유 있게 들어가는 **n(기본)과 s(선택)**를 씁니다.

Construction-PPE 데이터셋

Ultralytics가 공개한 공사 현장 안전장비 데이터셋입니다. 실제 현장 이미지 1,416장(학습 1,132 / 검증 143 / 테스트 141)에 11개 클래스가 라벨링돼 있습니다.

구분클래스
착용 장비 (5)helmet, gloves, vest, boots, goggles
미착용 (4)no_helmet, no_gloves, no_boots, no_goggle
기타 (2)Person, none

"착용"과 "미착용"을 따로 학습하기 때문에 단순 탐지를 넘어 안전 위반 감지까지 활용할 수 있습니다. 학습 코드에서 데이터셋 이름만 지정하면 자동으로 내려받습니다.

2. gcube 워크로드 만들기

gcube에서는 컨테이너 이미지 주소와 실행 명령을 직접 입력해 워크로드를 등록합니다. 아래 설정을 그대로 쓰면 RTX 5060에 맞는 PyTorch와 JupyterLab이 준비된 상태로 시작합니다.

항목입력값
GPURTX 5060 (8GB) × 1
저장소 유형Docker Hub
이미지pytorch/pytorch:2.9.1-cuda12.8-cudnn9-runtime
컨테이너 포트8888
최소 CUDA 버전12.8 이상

컨테이너 명령

bash -c "apt-get update && DEBIAN_FRONTEND=noninteractive apt-get install -y git wget libgl1 libglib2.0-0 && pip install jupyterlab && jupyter lab --ip=0.0.0.0 --port=8888 --no-browser --allow-root --ServerApp.token='' --ServerApp.password=''"

왜 이 이미지인가요? RTX 5060은 Blackwell 세대 GPU라서 CUDA 12.8 이상으로 빌드된 PyTorch가 필요합니다. 더 낮은 버전 이미지를 쓰면 GPU가 인식돼도 연산이 제대로 동작하지 않습니다. 또 libgl1을 미리 설치해 두어야 OpenCV 오류(libGL.so.1) 없이 바로 진행됩니다.

워크로드가 실행되면 서비스 주소 뒤에 /lab을 붙여 접속하고, Launcher에서 Notebook → Python 3을 엽니다. 토큰 없이 열리는 설정이므로 주소는 외부에 공유하지 마세요.

3. GPU 환경 확인

노트북 첫 셀에서 PyTorch가 RTX 5060을 제대로 인식하는지 확인합니다.

import torch
print("PyTorch:", torch.__version__)
print("CUDA:", torch.version.cuda)
print("GPU 사용 가능:", torch.cuda.is_available())
print("GPU 이름:", torch.cuda.get_device_name(0))
print("Compute capability:", torch.cuda.get_device_capability(0))
PyTorch: 2.9.1+cu128
CUDA: 12.8
GPU 사용 가능: True
GPU 이름: NVIDIA GeForce RTX 5060
Compute capability: (12, 0)

CUDA가 12.8 이상이고 sm_120 is not compatible 경고가 없으면 정상입니다.

4. Ultralytics 설치

!pip install -q ultralytics
import ultralytics
ultralytics.checks()
Ultralytics 8.4.178 🚀 Python-3.11.14 torch-2.9.1+cu128 CUDA:0 (NVIDIA GeForce RTX 5060, 8151MiB)
Setup complete ✅ (12 CPUs, 20.5 GB RAM, 60.8/75.6 GB disk)

Setup complete ✅가 보이면 준비 끝입니다.

5. 첫 실행: 사전학습 모델로 바로 탐지해 보기

학습에 들어가기 전에, 사람·차·버스 등 80종 사물을 이미 알고 있는 기본 모델이 잘 도는지 확인합니다.

from ultralytics import YOLO
from ultralytics.utils import ASSETS
from IPython.display import Image, display

model = YOLO("yolo26n.pt")  # 처음 실행 시 자동 다운로드
results = model.predict(ASSETS / "bus.jpg", save=True)  # 패키지 내장 예제 이미지

save_dir = results[0].save_dir
print("결과 저장 위치:", save_dir)
display(Image(filename=f"{save_dir}/bus.jpg", width=600))

실제 출력은 640x480 4 persons, 1 bus, 6.1ms였습니다. 사진 한 장을 분석하는 데 0.006초 정도 걸린 셈입니다.

6. 안전장비 탐지 모델 학습

6-1. YOLO26n 학습 (권장)

import time, torch
from ultralytics import YOLO

torch.cuda.reset_peak_memory_stats()
t0 = time.time()
model = YOLO("yolo26n.pt")
results = model.train(
    data="construction-ppe.yaml",  # 데이터셋 자동 다운로드
    epochs=50,
    imgsz=640,
    batch=16,
    device=0,
    workers=4,
    project="ppe",
    name="y26n",
)
print(f"학습 시간: {(time.time()-t0)/60:.1f}분")
print(f"최대 VRAM: {torch.cuda.max_memory_allocated()/1024**3:.2f} GB")

실행하면 데이터셋(170MB)을 받은 뒤 바로 학습이 시작됩니다. 로그의 GPU_mem을 보면 8GB 중 3.2GB 정도만 쓰는 것을 확인할 수 있습니다.

옵션의미8GB 팁
epochs전체 데이터를 몇 번 반복 학습할지50으로 충분
batch한 번에 처리하는 이미지 수메모리 부족 시 8로 낮춤
workers데이터 로딩 프로세스 수공유 메모리 오류 시 2 또는 0

6-2. YOLO26s로 더 높은 정확도 (선택)

6-1 직후 같은 커널에서 실행하면 앞 학습의 GPU 메모리가 남아 CUDA error: out of memory가 납니다. 메뉴 Kernel → Restart Kernel… → Restart로 커널을 재시작한 뒤 실행하세요.

import time, torch
from ultralytics import YOLO

torch.cuda.reset_peak_memory_stats()
t0 = time.time()
model = YOLO("yolo26s.pt")
results = model.train(data="construction-ppe.yaml", epochs=100, imgsz=640, batch=16,
                      device=0, workers=4, project="ppe", name="y26s")
print(f"학습 시간: {(time.time()-t0)/60:.1f}분")
print(f"최대 VRAM: {torch.cuda.max_memory_allocated()/1024**3:.2f} GB")

7. 학습 결과 확인

from IPython.display import Image, display  # 커널을 재시작했다면 다시 불러와야 함
save_dir = model.trainer.save_dir  # 실제 저장 위치
print("학습 결과 폴더:", save_dir)
display(Image(filename=f"{save_dir}/results.png", width=900))
display(Image(filename=f"{save_dir}/val_batch0_pred.jpg", width=900))

results.png는 epoch별 손실과 정확도 그래프, val_batch0_pred.jpg는 검증 이미지 예측 예시입니다. 학습하지 않은 테스트셋으로 최종 성능도 측정합니다.

best = YOLO(f"{save_dir}/weights/best.pt")
metrics = best.val(data="construction-ppe.yaml", split="test")
print("mAP50:", round(metrics.box.map50, 3))
print("mAP50-95:", round(metrics.box.map, 3))

실측 결과 (RTX 5060 8GB, batch 16)

모델epoch학습 시간최대 VRAM테스트 mAP50테스트 mAP50-95
YOLO26n507.5분2.6GB0.5350.267
YOLO26s10021.7분4.6GB0.5620.273
  • 두 모델 모두 8GB의 절반 안팎으로 여유 있게 학습됐습니다.
  • s는 n보다 정확도가 조금 높지만 시간이 약 3배 걸립니다. 처음에는 n으로 흐름을 익히는 것을 권장합니다.
  • 클래스별로는 헬멧·조끼 같은 착용 장비는 잘 찾지만, 미착용 클래스는 정확도가 낮았습니다. 미착용 사례가 학습 데이터에 적기 때문입니다.
  • 같은 이미지·설정으로 두 번 학습해도 수치가 똑같이 나왔습니다. 따라 하면 거의 같은 결과를 얻을 수 있습니다.

mAP 읽는 법: 0~1 사이 값으로, 모델이 박스를 얼마나 정확히 찾았는지 나타냅니다. mAP50은 "대략 맞으면 정답", mAP50-95는 "위치까지 정밀해야 정답" 기준입니다.

8. 내 사진에 적용하기

JupyterLab 왼쪽 파일 패널에 현장 사진(예: site.jpg)을 업로드한 뒤 실행합니다.

from PIL import Image as PILImage

res = best.predict("site.jpg", conf=0.4, save=True)
print([(best.names[int(b.cls)], round(float(b.conf), 2)) for b in res[0].boxes])

display(PILImage.fromarray(res[0].plot()[..., ::-1]))  # 박스가 그려진 결과 바로 보기

테스트셋 현장 사진으로 실행한 결과는 [('helmet', 0.82), ('vest', 0.81), ('Person', 0.78), ('gloves', 0.51)]였습니다. 헬멧·조끼·장갑을 착용한 작업자를 찾아낸 것입니다.

  • conf=0.4는 확신도 40% 이상만 표시한다는 뜻입니다. 오탐이 많으면 올리고, 놓치는 게 많으면 내립니다.
  • 영상도 best.predict("site.mp4", save=True)로 같은 방식으로 처리합니다.

미착용만 골라내기 — 안전 관제 시스템의 기본 형태입니다.

VIOLATIONS = {"no_helmet", "no_gloves", "no_boots", "no_goggle"}
found = [best.names[int(b.cls)] for b in res[0].boxes if best.names[int(b.cls)] in VIOLATIONS]
print("위반 감지:", found if found else "없음")

9. 모델 내보내기 (ONNX)

다른 서버나 엣지 기기에서 쓰려면 ONNX 형식으로 변환합니다.

onnx_path = best.export(format="onnx")
print(onnx_path)

처음 실행하면 필요한 패키지가 자동 설치되고 몇 초 안에 best.onnx가 만들어집니다. 중간에 Restart runtime... 경고가 보여도 Export complete가 나오면 성공입니다. 파일 패널에서 우클릭 → Download로 받을 수 있습니다.

10. 트러블슈팅

아래는 다른 이미지를 쓰거나 순서를 바꿨을 때 실제로 겪은 오류입니다.

증상원인해결
sm_120 is not compatible 경고PyTorch가 CUDA 12.4 이하 빌드2단계 권장 이미지 사용, 또는 pip install --upgrade torch torchvision --index-url https://download.pytorch.org/whl/cu128 후 커널 재시작
ImportError: libGL.so.1컨테이너에 OpenCV용 라이브러리 없음2단계 명령에 libgl1 포함, 또는 opencv-python-headless로 교체 후 커널 재시작
Download failure ... Connection reset by peer다운로드 서버 간헐적 연결 끊김셀 재실행. 데이터셋은 아래 코드로 직접 받기
CUDA error: out of memory앞 학습의 GPU 메모리가 남아 있음커널 재시작 후 실행, 그래도 부족하면 batch=8

데이터셋 자동 다운로드가 실패할 때:

from ultralytics.utils import DATASETS_DIR
from ultralytics.utils.downloads import download

download("https://github.com/ultralytics/assets/releases/download/v0.0.0/construction-ppe.zip", dir=DATASETS_DIR)

마치며

8GB GPU 한 장으로 30분 만에 공사 현장 안전장비를 찾는 모델을 만들었습니다. 학습에는 VRAM 2.6GB만 썼으니, 객체 탐지 같은 비전 AI는 고가의 GPU 없이도 충분히 시작할 수 있습니다.

다음으로 해 볼 것

  • 나만의 데이터로 학습: Label Studio나 CVAT로 YOLO 형식 라벨을 만들고, 같은 구조의 yaml 파일만 작성하면 같은 코드로 학습할 수 있습니다.
  • 다른 작업으로 확장: 같은 API로 세그멘테이션(yolo26n-seg.pt), 자세 추정(yolo26n-pose.pt)도 가능합니다.
  • 더 큰 모델: 더 높은 정확도가 필요하면 gcube의 더 큰 GPU에서 m·l 모델을 학습해 보세요.

라이선스 안내: Ultralytics YOLO 코드와 Construction-PPE 데이터셋은 AGPL-3.0 라이선스입니다. 학습·연구 목적으로는 자유롭게 쓸 수 있지만, 이를 포함한 서비스를 상용 배포할 때는 소스 공개 의무가 생기므로 Ultralytics 엔터프라이즈 라이선스를 검토해야 합니다.

참고 자료