Pipeline 성능 — Pin Memory, Workers, Prefetch

GPU 가 80% 활용도 아래면 pipeline 이 bottleneck

완벽히 먹은 GPU 가 training 중 100% 가까이. 이하면 data pipeline 이 compute 를 바닥에 남기는 중. DataLoader 에 fix knob 있어 — 대부분 한 번 설정 후 잊음.

중요한 setting

num_workers — 병렬 __getitem__ 위 subprocess 수. CPU core 수로 시작; 메모리 빡빡하면 down-tune (각 worker 가 dataset copy). 0 = main-process loading (느리지만 디버그 쉬움).
pin_memory=True — non-pageable CPU memory 에 batch 할당. x.to(device, non_blocking=True) 와 결합 시 CPU↔GPU overlap 가능.
prefetch_factor — 각 worker 가 미리 prefetch 하는 batch. default 2; GPU 가 batch 보다 빨리 끝내면 4 로.
persistent_workers=True — epoch 들 사이 worker 살아있게. 첫 epoch 이 worker-startup 세금 지불; persistent 가 전체 training run 에 한 번.
drop_last=True — 뒤 partial batch drop. 일관된 batch shape (BatchNorm stats, distributed training) 위 유용.

bottleneck 진단

nvidia-smi 가 GPU 30%? Pipeline 이 bottleneck.
nvidia-smi 가 GPU 95%? Pipeline 괜찮음; speedup 은 compute 에서 (compile, AMP, 더 큰 batch).
htop 이 한 CPU core 100%, 다른 거 idle? num_workers 너무 낮음.
multi-worker loading 중 메모리 압력? 각 worker 가 dataset copy; num_workers 낮추거나 IterableDataset.

Code

Production-shaped 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...

dataset normalization stats 계산·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]

Quick benchmark — loader 충분히 빠름?·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)

Pipeline 성능 — Pin Memory, Workers, Prefetch

GPU 가 80% 활용도 아래면 pipeline 이 bottleneck

중요한 setting

bottleneck 진단

Code

External links

Exercise

Progress

댓글 0