본문 바로가기
Artificial Intelligence/60. Python

[PYTHON] DataLoader num_workers 설정이 학습 속도와 메모리에 미치는 3가지 영향과 해결 방법

by Papa Martino V 2026. 4. 18.
728x90

DataLoader num_workers 설정
DataLoader num_workers 설정

 

딥러닝 모델 학습 중 GPU 사용률(GPU Utilization)이 낮게 유지되거나 학습 속도가 기대보다 느리다면, 범인은 모델 아키텍처가 아닌 데이터 로딩 파이프라인일 확률이 매우 높습니다. 파이썬(Python) 기반의 PyTorch 프레임워크에서 DataLoadernum_workers 옵션은 데이터를 읽어오는 멀티프로세싱의 수준을 결정하는 핵심 스위치입니다. 본 포스팅에서는 num_workers가 시스템 자원에 미치는 구조적 차이를 심층 분석하고, 데이터 로딩 병목을 해결하여 학습 효율을 200% 이상 끌어올릴 수 있는 7가지 실전 테크닉을 공유합니다.


1. num_workers 수치에 따른 시스템 동작 차이와 병목 지점

num_workers는 메인 프로세스가 데이터를 기다리는 동안, 별도의 서브 프로세스들이 데이터를 미리 로드하고 전처리하도록 합니다. 이 수치를 높이면 속도는 빨라지지만, 메모리(RAM) 점유율과 CPU 부하가 기하급수적으로 증가하는 트레이드오프(Trade-off)가 발생합니다.

설정 값 동작 메커니즘 장점 (Pros) 단점 및 위험 요소 (Cons)
0 (기본값) 메인 프로세스에서 동기식 로드 디버깅 용이, 메모리 절약 심각한 GPU 대기 병목 발생
1 ~ CPU 코어 절반 멀티프로세싱 병렬 로드 학습 속도 대폭 향상 프로세스 간 IPC 오버헤드 발생
CPU 코어 초과 과도한 컨텍스트 스위칭 이론상 최대 병렬화 Shared Memory 부족 및 속도 저하
최적 수치 GPU-CPU 밸런스 최적화 최고의 Throughput 달성 환경별 벤치마킹 필요

2. 실무 데이터 파이프라인 최적화를 위한 7가지 해결 방법 (Examples)

다음은 현업에서 대규모 데이터셋을 다룰 때 num_workers와 관련하여 발생하는 문제를 해결하고 성능을 극대화하는 파이썬 코드 예시입니다.

Example 1: 현재 하드웨어 사양에 맞춘 자동 num_workers 할당

코드가 실행되는 서버의 CPU 코어 수를 감지하여 적절한 워커 수를 동적으로 설정하는 해결 방법입니다.

import os
import torch
from torch.utils.data import DataLoader

def get_optimal_workers():
    # CPU 코어 수의 약 4배 혹은 사용 가능한 코어 수로 설정 (서버 환경에 따라 조절)
    num_cores = os.cpu_count()
    return num_cores if num_cores is not None else 0

loader = DataLoader(
    dataset, 
    batch_size=64, 
    num_workers=get_optimal_workers(),
    pin_memory=True
)

Example 2: pin_memory를 활용한 CPU-GPU 전송 가속

워커가 로드한 데이터를 GPU 메모리로 더 빠르게 복사할 수 있도록 페이지 고정 메모리(Page-locked memory)를 사용하는 방법입니다.

# num_workers가 0보다 클 때 pin_memory=True는 필수적입니다.
train_loader = DataLoader(
    dataset, 
    batch_size=128, 
    num_workers=4, 
    pin_memory=True, # 메모리 전송 병목 해결
    shuffle=True
)

Example 3: prefetch_factor 설정을 통한 워커 유휴 시간 방지

각 워커가 배치를 미리 몇 개나 준비해둘지 결정하여 CPU의 연산 자원을 쉬지 않게 만드는 기법입니다.

# PyTorch 1.7+ 버전부터 지원
loader = DataLoader(
    dataset, 
    batch_size=32, 
    num_workers=8,
    prefetch_factor=3, # 각 워커당 3개의 배치를 미리 로드
    persistent_workers=True # 에포크 종료 시 워커 유지
)

Example 4: persistent_workers로 에포크 전환 시 지연 현상 제거

새 에포크가 시작될 때마다 워커를 새로 생성하는 오버헤드를 줄여주는 설정입니다.

# 데이터 로딩 초기화 시간이 길 때 매우 효과적임
loader = DataLoader(
    dataset,
    num_workers=4,
    persistent_workers=True # 프로세스를 죽이지 않고 다음 에포크에서 재사용
)

Example 5: 데이터 로딩 속도 측정을 위한 프로파일링 코드

현재 설정한 num_workers가 실제로 효율적인지 확인하기 위해 배치 로딩 시간을 측정하는 유틸리티입니다.

import time

def profile_loader(loader):
    start_time = time.time()
    for i, data in enumerate(loader):
        if i > 10: break # 처음 10개 배치만 테스트
        pass
    end_time = time.time()
    print(f"평균 배치 로딩 시간: {(end_time - start_time) / 10:.4f}s")

# 다른 num_workers 값으로 성능을 비교해 보세요.

Example 6: Shared Memory(shm) 부족 에러 해결을 위한 설정

도커(Docker) 환경에서 num_workers를 높일 때 발생하는 공유 메모리 부족 문제를 방지하는 전략입니다.

# Python 코드가 아닌 Docker 실행 시 해결 방법 (가이드)
# --shm-size=2g 와 같이 충분한 공유 메모리를 할당해야 합니다.
# 만약 할당이 어렵다면, num_workers를 줄이거나 데이터를 압축 처리하세요.

Example 7: 가변 길이 데이터를 위한 맞춤형 Collate Function 최적화

워커들이 배치를 묶는 과정에서 발생하는 CPU 연산을 최소화하여 병목을 줄이는 해결 예시입니다.

def efficient_collate_fn(batch):
    # 불필요한 복사를 피하고 텐서 변환을 최소화하는 로직 구현
    images = [item[0] for item in batch]
    labels = [item[1] for item in batch]
    return torch.stack(images, dim=0), torch.tensor(labels)

loader = DataLoader(dataset, collate_fn=efficient_collate_fn, num_workers=4)

3. 결론: 성능과 안정성 사이의 골디락스 존 찾기

num_workers 설정의 핵심은 "GPU가 데이터를 기다리게 하지 않으면서도, CPU와 RAM에 과부하를 주지 않는 지점"을 찾는 것입니다. 일반적으로 사용 가능한 CPU 코어 수의 2배에서 4배 사이를 시작점으로 하되, pin_memoryprefetch_factor를 조합하여 최적의 Throughput을 찾아야 합니다. 특히 고해상도 이미지 처리가 포함된 멀티모달 학습에서는 워커 하나당 점유하는 메모리가 크므로, 무작정 워커 수를 늘리기보다 데이터 직렬화(TFRecord, WebDataset)를 먼저 고려하는 것이 현명한 접근입니다.


4. 전문 지식 출처 및 참고 문헌

  • PyTorch Documentation: "Data Loader Workers and Memory Pinning"
  • NVIDIA Developer Blog: "Advanced PyTorch Data Loading Techniques"
  • Academic Study: "Empirical Analysis of Multi-processing Data Loading in Deep Learning Frameworks"
728x90