[Ultralytics] 학습 중 Multi-GPU 검증 지원

PR 링크: ultralytics/ultralytics#22377 상태: Merged | 변경: +209 / -59

들어가며

Ultralytics YOLO의 Multi-GPU(DDP) 학습에는 오랜 병목이 있었다. 학습은 모든 GPU에서 병렬로 수행되지만, 에폭 종료 후 검증(validation)은 Rank 0 GPU 한 대에서만 실행되었다. GPU가 8장이면 7장은 검증 동안 놀고 있었다는 뜻이다. 이 PR은 검증 단계도 모든 GPU에서 병렬 수행하도록 변경하여, 학습 전체 시간을 단축한다.

핵심 코드 분석

1. ContiguousDistributedSampler 도입

기존 PyTorch의 DistributedSampler는 round-robin 방식으로 데이터를 분배한다(GPU 0: [0,2,4,...], GPU 1: [1,3,5,...]). 이러면 rect=True처럼 이미지 크기별로 정렬된 데이터셋의 순서가 깨진다.

Before:

sampler = None if rank == -1 else distributed.DistributedSampler(dataset, shuffle=shuffle)

After:

sampler = (
    None
    if rank == -1
    else distributed.DistributedSampler(dataset, shuffle=shuffle)
    if shuffle
    else ContiguousDistributedSampler(dataset)
)

검증(shuffle=False)일 때는 새로운 ContiguousDistributedSampler를 사용한다. 이 sampler는 각 GPU에 연속된 배치 블록을 할당하여 데이터셋의 정렬 순서를 유지한다.

2. 검증 로직을 모든 Rank로 확장

Before:

if RANK in {-1, 0}:
    self.test_loader = self.get_dataloader(
        self.data.get("val") or self.data.get("test"),
        batch_size=batch_size if self.args.task == "obb" else batch_size * 2,
        rank=-1,
        mode="val",
    )
    self.validator = self.get_validator()

After:

self.test_loader = self.get_dataloader(
    self.data.get("val") or self.data.get("test"),
    batch_size=batch_size if self.args.task == "obb" else batch_size * 2,
    rank=LOCAL_RANK,
    mode="val",
)
self.validator = self.get_validator()
self.ema = ModelEMA(self.model)

test_loader와 validator 생성을 if RANK in {-1, 0} 블록 바깥으로 이동시켜, 모든 GPU에서 검증 인프라를 초기화한다.

3. 통계 수집 및 Loss 동기화

검증이 각 GPU에서 부분 데이터로 수행되므로, 결과를 Rank 0에 모아야 한다.

After (validator.py):

self.gather_stats()
if RANK in {-1, 0}:
    stats = self.get_stats()
    self.speed = dict(zip(self.speed.keys(), (x.t / len(self.dataloader.dataset) * 1e3 for x in dt)))
    self.finalize_metrics()
    self.print_results()

After (detect/val.py):

def gather_stats(self) -> None:
    if RANK == 0:
        gathered_stats = [None] * dist.get_world_size()
        dist.gather_object(self.metrics.stats, gathered_stats, dst=0)
        merged_stats = {key: [] for key in self.metrics.stats.keys()}
        for stats_dict in gathered_stats:
            for key in merged_stats.keys():
                merged_stats[key].extend(stats_dict[key])
        self.metrics.stats = merged_stats
    elif RANK > 0:
        dist.gather_object(self.metrics.stats, None, dst=0)

dist.gather_object로 각 GPU의 통계를 Rank 0에 수집한 뒤 병합한다. Loss도 dist.reduce로 평균을 계산한다.

왜 이게 좋은가

• 검증 시간 N분의 1: GPU 8장이면 검증 데이터를 8등분하여 병렬 처리한다
• 데이터 정렬 유지: ContiguousDistributedSampler가 연속 블록을 할당하므로 rect=True 최적화가 깨지지 않는다
• EMA 동기화: 검증 전 dist.broadcast로 EMA 버퍼를 Rank 0에서 전파하여 모든 GPU가 동일한 모델로 검증한다

정리

DDP 학습의 검증 병목을 해소한 실용적인 PR이다. 핵심 아이디어는 세 가지다: (1) 연속 블록 분배 sampler, (2) 검증 결과의 분산 수집/병합, (3) EMA 버퍼 브로드캐스트. Multi-GPU 학습 파이프라인을 설계할 때 검증 단계도 반드시 병렬화를 고려해야 한다는 교훈을 준다.

참고 자료

• PyTorch DistributedSampler 문서 — 기본 분산 sampler 동작 원리
• Ultralytics DDP Training Guide — Ultralytics Multi-GPU 학습 가이드

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.