파이프라인 성능: 페이지 고정 메모리, 작업 프로세스, 미리 불러오기

GPU 활용도가 80% 아래라면 데이터 파이프라인이 병목일 수 있어

계산할 일을 충분히 받은 GPU는 학습 중 활용도가 100%에 가까워. 그보다 크게 낮다면 GPU가 데이터를 기다리며 쉬고 있을 가능성이 커. DataLoader의 몇 가지 설정으로 고칠 수 있고, 보통은 한 번 맞춘 뒤 그대로 사용해.

중요한 설정

num_workers: __getitem__을 병렬로 실행할 작업 프로세스 수야. CPU 코어 수를 기준으로 시작하되, 각 작업 프로세스가 데이터셋을 복제하므로 메모리가 빠듯하면 낮춰. 0이면 주 프로세스에서 불러와 느리지만 디버깅은 쉬워.
pin_memory=True: 배치를 페이지 고정 CPU 메모리에 할당해. x.to(device, non_blocking=True)와 함께 쓰면 CPU→GPU 전송과 계산을 겹칠 수 있어.
prefetch_factor: 각 작업 프로세스가 미리 불러올 배치 수야. 기본값은 2고, GPU가 다음 배치보다 먼저 계산을 끝낸다면 4로 올려 봐.
persistent_workers=True: 에포크 사이에도 작업 프로세스를 유지해. 첫 에포크에서만 시작 비용을 내고 학습 내내 재사용할 수 있어.
drop_last=True: 마지막 자투리 배치를 버려. BatchNorm 통계나 분산 학습처럼 배치 모양을 일정하게 유지해야 할 때 유용해.

병목 진단

nvidia-smi에서 GPU 활용도가 30%라면 데이터 파이프라인이 병목일 가능성이 커.
GPU 활용도가 95%라면 파이프라인은 괜찮아. 컴파일, AMP, 더 큰 배치처럼 계산 쪽에서 속도를 높여야 해.
htop에서 CPU 코어 하나만 100%이고 나머지가 놀고 있다면 num_workers가 너무 낮아.
여러 작업 프로세스로 불러올 때 메모리 압력이 크다면 데이터셋 복제 비용이 원인일 수 있어. num_workers를 낮추거나 IterableDataset을 고려해.

Code

운영 환경을 고려한 DataLoader·python

import torch
from torch.utils.data import DataLoader

loader = DataLoader(
    dataset,
    batch_size=64,
    shuffle=True,
    num_workers=8,                # try matching CPU cores
    pin_memory=True,              # fast CPU→GPU
    prefetch_factor=4,            # extra prefetch headroom
    persistent_workers=True,      # don't re-spin per epoch
    drop_last=True,               # consistent batch shape
)

device = "cuda"
for x, y in loader:
    x = x.to(device, non_blocking=True)
    y = y.to(device, non_blocking=True)
    # ...training...

데이터셋 정규화 통계 계산·python

import torch
from torch.utils.data import DataLoader

def compute_mean_std(loader):
    """Compute per-channel mean and std across an image dataset."""
    n_pixels = 0
    sum_  = torch.zeros(3)
    sum_sq = torch.zeros(3)
    for x, _ in loader:
        b, c, h, w = x.shape
        n_pixels += b * h * w
        sum_  += x.sum(dim=(0, 2, 3))
        sum_sq += (x ** 2).sum(dim=(0, 2, 3))
    mean = sum_ / n_pixels
    var  = sum_sq / n_pixels - mean ** 2
    std  = torch.sqrt(var)
    return mean, std

# ImageNet defaults (memorize these): mean=[.485,.456,.406] std=[.229,.224,.225]

빠른 성능 측정: 로더 충분히 빠름?·python

import time
import torch
from torch.utils.data import DataLoader

def benchmark_loader(loader, n=100, device='cuda'):
    it = iter(loader)
    t0 = time.perf_counter()
    samples = 0
    for i in range(n):
        x, _ = next(it)
        x = x.to(device, non_blocking=True)
        samples += x.size(0)
    if device == 'cuda':
        torch.cuda.synchronize()
    elapsed = time.perf_counter() - t0
    print(f"{samples/elapsed:,.0f} samples/sec")

# Compare configurations
for n_workers in [0, 2, 4, 8]:
    loader = DataLoader(dataset, batch_size=64, num_workers=n_workers,
                        pin_memory=True, persistent_workers=(n_workers > 0))
    benchmark_loader(loader)

Exercise

세 번째 코드 블록의 성능 측정을 원하는 데이터셋에서 num_workers를 0, 2, 4, 8로 바꿔 실행해 봐. 초당 샘플 수와 num_workers의 관계를 그래프로 그리면 보통 처음에는 가파르게 오르다가 어느 지점부터 정체돼. 그 지점의 값을 실제 실행에 사용해.