PyTorch DataLoader 최적화: CPU 과부하 없이 효율적인 데이터 처리하기

딥러닝 모델 학습 시 데이터 로딩이 병목이 되어 GPU가 놀고 있는 경험, 한 번쯤은 있으실 겁니다. 특히 num_workers를 늘렸는데 오히려 CPU 사용률만 치솟고 학습 속도는 개선되지 않는 상황이 발생하곤 합니다. 이 글에서는 PyTorch DataLoader를 효율적으로 사용하는 방법과 CPU 자원을 최적화하는 전략을 상세히 알아보겠습니다.

DataLoader에서 NumPy vs PyTorch: 언제 무엇을 써야 할까?

NumPy가 유리한 경우 (일반적인 상황)

데이터 로딩과 전처리는 주로 CPU에서 이루어지므로, 대부분의 경우 NumPy가 더 자연스럽고 효율적입니다.

NumPy의 장점:

• C-level 최적화: 내부적으로 C, Fortran으로 작성되어 CPU에서 매우 빠르게 동작
• 메모리 효율성: 연속적인 메모리 블록에 저장되어 메모리 접근이 효율적
• 라이브러리 호환성: PIL, OpenCV 등 이미지 처리 라이브러리와 호환성이 좋음

DataLoader의 워커 프로세스는 NumPy 배열로 전처리한 후 torch.from_numpy()를 통해 PyTorch 텐서로 변환하여 메인 프로세스로 전달합니다. 이 과정이 CPU에서 독립적으로 병렬 처리되므로 매우 효율적입니다.

PyTorch가 유리한 경우 (특수한 상황)

GPU 전처리 가능한 경우:

• 전처리 작업이 간단하고 텐서 연산만으로 충분하다면
• NumPy 대신 PyTorch 텐서로 데이터를 읽은 후 GPU로 전처리를 옮기는 것이 가장 효율적
• NVIDIA DALI 같은 라이브러리가 이 원리를 활용

NumPy 사용 시 CPU 과부하 발생하는 경우:

• NumPy는 내부적으로 멀티스레딩을 활용하므로
• num_workers를 높게 설정하면 각 워커가 여러 스레드를 생성하여 CPU 사용률 급증
• 이 경우 NumPy 대신 PyTorch를 사용하거나 스레드 개수를 제한해야 함

torch.set_num_threads() 완벽 가이드

왜 필요한가?

torch.set_num_threads() 함수는 PyTorch 연산에 사용될 CPU 스레드의 최대 개수를 설정합니다. 워커 프로세스가 불필요하게 많은 스레드를 생성하여 CPU 과부하를 일으키는 것을 방지하기 위해 사용합니다.

올바른 사용법

방법 1: 메인 스크립트 상단에 설정

import torch
import os
from torch.utils.data import DataLoader, Dataset
import numpy as np

# PyTorch가 사용하는 CPU 스레드 수를 1개로 제한
torch.set_num_threads(1)

# NumPy, OpenMP 등에서 사용하는 스레드 개수도 함께 제한
os.environ['OMP_NUM_THREADS'] = '1'
os.environ['MKL_NUM_THREADS'] = '1'
os.environ['OPENBLAS_NUM_THREADS'] = '1'

# 데이터셋 정의
class MyDataset(Dataset):
    def __init__(self):
        self.data = np.random.rand(10000, 3, 224, 224).astype(np.float32)
        self.labels = np.random.randint(0, 10, size=10000)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        # NumPy 연산 (예시: 이미지 전처리)
        processed_data = self.data[idx] * 255.0
        
        # NumPy 배열을 PyTorch 텐서로 변환
        return torch.from_numpy(processed_data), self.labels[idx]

dataset = MyDataset()
dataloader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4)

# 학습 루프
for batch_idx, (data, target) in enumerate(dataloader):
    data = data.to('cuda')
    target = target.to('cuda')
    # ... 모델 학습 코드

핵심 포인트

• 위치가 중요: DataLoader가 생성되기 전, 즉 스크립트의 최상단에 위치
• __getitem__ 메서드 내부에 넣으면 안 되는 이유: 매 샘플을 불러올 때마다 함수가 호출되어 성능 저하 유발
• 워커 프로세스 상속: 메인 프로세스의 설정이 모든 워커 프로세스에 상속됨

CPU vs GPU: 역할 분담이 핵심

CPU 과부하를 예방하는 올바른 접근법

1. 스레드 수 제한하기

# 환경 변수 설정으로 각 워커가 사용할 스레드 수 제한
os.environ['OMP_NUM_THREADS'] = '1'
os.environ['MKL_NUM_THREADS'] = '1'
os.environ['OPENBLAS_NUM_THREADS'] = '1'
torch.set_num_threads(1)

이렇게 설정하면:

• 총 워커 프로세스 수: num_workers개
• 각 워커 프로세스가 사용하는 스레드 수: 1개
• 전체 CPU 코어 사용: 최대 num_workers개의 코어만 사용

DataLoader에서 CUDA 변환하면 안 되는 이유

많은 분들이 "DataLoader 안에서 텐서를 cuda로 변환해서 연산하면 되지 않을까?"라고 생각하시는데, 이는 권장되지 않습니다.

문제점들:
1. 리소스 충돌: 여러 워커 프로세스가 동시에 GPU에 접근하려고 시도
2. GPU 메모리 오버헤드: 각 워커가 GPU 메모리를 할당받고 해제하는 오버헤드 발생
3. 병목 현상: DataLoader의 본래 목적인 CPU에서의 빠른 데이터 전처리가 무의미해짐

최적화된 워크플로우

권장하는 완벽한 파이프라인

import os
import torch
from torch.utils.data import DataLoader, Dataset
import torchvision.transforms as transforms
import numpy as np

# 1. 스크립트 상단에서 스레드 수 제한 (가장 중요!)
os.environ['OMP_NUM_THREADS'] = '1'
os.environ['MKL_NUM_THREADS'] = '1' 
os.environ['OPENBLAS_NUM_THREADS'] = '1'
torch.set_num_threads(1)

# 2. 데이터셋 정의 (CPU에서 전처리)
class OptimizedDataset(Dataset):
    def __init__(self, data, labels, transform=None):
        self.data = data
        self.labels = labels
        self.transform = transform

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        # NumPy에서 데이터 로딩
        image = self.data[idx]
        label = self.labels[idx]

        # CPU에서 전처리
        if self.transform:
            image = self.transform(image)
            
        return image, label

# 3. DataLoader 생성
train_dataset = OptimizedDataset(data, labels, transform=transforms.ToTensor())
train_loader = DataLoader(
    train_dataset, 
    batch_size=32, 
    shuffle=True, 
    num_workers=4,  # CPU 코어 수에 맞게 조정
    pin_memory=True  # GPU 복사 속도 향상
)

# 4. 메인 학습 루프 (여기서 GPU로 이동)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

for batch_idx, (data, target) in enumerate(train_loader):
    # GPU로 데이터 이동은 메인 루프에서!
    data, target = data.to(device), target.to(device)
    
    # GPU에서 추가 전처리 가능
    # data = some_gpu_augmentation(data)
    
    # 모델 학습
    # ...

핵심 정리

CPU 자원을 최대한 아끼는 방법

1. 스레드 제한: torch.set_num_threads(1) 및 환경 변수 설정
2. 역할 분담: CPU는 데이터 로딩, GPU는 연산
3. 적절한 num_workers: CPU 코어 수에 맞게 설정
4. pin_memory=True: CPU-GPU 간 데이터 전송 속도 향상

주의사항

• ❌ DataLoader 내부에서 .to('cuda') 사용하면 안됨
• ❌ __getitem__ 메서드 내부에 torch.set_num_threads() 배치하면 안됨
• ❌ num_workers를 너무 크게 설정하면 안됨
• ✅ 스크립트 최상단에 스레드 제한 설정하기
• ✅ CPU-GPU 역할 명확히 분리하기
• ✅ 시스템 리소스 모니터링으로 최적값 찾기하기

효율적인 데이터 로딩은 딥러닝 성능의 핵심입니다. 이 가이드를 따라하시면 CPU 과부하 없이 GPU를 최대한 활용하는 최적화된 학습 환경을 구축하실 수 있을 것입니다.

---