BioLab YOLO Training

Backend.AI (SKKU Supercomputing Center) GPU 클러스터에서 YOLO 모델을 학습하기 위한 스크립트입니다.

파일 구조

model_train/
├── train.py              # 메인 학습 스크립트
├── requirements.txt      # Python 의존성
└── README.md            # 이 파일

사용 방법

1. Backend.AI에서 실행

Step 1: 데이터셋 업로드

VFolder에 다음 구조로 데이터셋 업로드:

data/
├── images/
│   ├── train/
│   │   ├── img001.jpg
│   │   └── ...
│   └── val/
│       ├── img050.jpg
│       └── ...
├── labels/
│   ├── train/
│   │   ├── img001.txt  (YOLO format)
│   │   └── ...
│   └── val/
│       └── ...
└── data.yaml

Step 2: data.yaml 예시

path: /home/work/data
train: images/train
val: images/val
names:
  0: pad
nc: 1

Step 3: Batch 세션 생성

• 이미지: pytorch:24.09-py310-cuda12.6 (또는 유사)
• 리소스: GPU 1개, RAM 32GB
• 마운트: 데이터셋 VFolder → /home/work/data
• 실행 명령어:

pip install -r requirements.txt && python train.py --data /home/work/data/data.yaml --epochs 100

2. 학습 파라미터

파라미터	기본값	설명
--data	(필수)	data.yaml 경로
--epochs	100	학습 에폭 수
--batch	16	배치 사이즈
--imgsz	640	이미지 크기
--model	yolov8n.pt	베이스 모델
--project	/home/work/output	출력 디렉토리
--name	train	실행 이름
--augment	True	데이터 증강

3. 출력 결과

학습 완료 후 /home/work/output/train/ 폴더에:

• weights/best.pt - 최고 성능 모델
• weights/last.pt - 마지막 에폭 모델
• results.png - 학습 그래프
• training_summary.json - 학습 요약

예시 명령어

# 기본 학습
python train.py --data /home/work/data/data.yaml

# 빠른 테스트 (10 에폭)
python train.py --data /home/work/data/data.yaml --epochs 10

# 큰 배치 사이즈 (A100 사용 시)
python train.py --data /home/work/data/data.yaml --batch 32

# 증강 없이 학습
python train.py --data /home/work/data/data.yaml --no-augment

주의사항

1. GPU 필수: 이 스크립트는 GPU가 필요합니다.
2. YOLO 형식: 라벨은 YOLO 형식 (클래스 cx cy w h)이어야 합니다.
3. 경로 주의: Backend.AI에서는 /home/work/가 기본 작업 디렉토리입니다.

문제 해결

CUDA out of memory

• --batch 값을 줄이세요 (16 → 8 → 4)

data.yaml not found

• VFolder 마운트 경로를 확인하세요
• --data 경로가 정확한지 확인하세요

학습이 너무 느림

• A100 GPU 사용을 권장합니다
• --imgsz 640 또는 --imgsz 480으로 줄여보세요