[PYTORCH] 거대한 데이터셋을 메모리 부족 없이 로드하는 7가지 전략 및 성능 해결 방법

Out of Memory (OOM)

 

현대 딥러닝 아키텍처는 수백 기가바이트에서 테라바이트 단위의 데이터를 필요로 합니다. 하지만 개발자의 워크스테이션이나 클라우드 인스턴스의 RAM 용량은 한계가 있습니다. PyTorch 환경에서 Out of Memory (OOM) 에러를 피하면서 수억 개의 샘플을 처리하려면 데이터셋 설계 단계부터 Lazy Loading과 Data Streaming 전략을 도입해야 합니다. 본 가이드에서는 2026년 실무 표준인 대규모 데이터 핸들링 기법을 상세히 다룹니다.

---

1. 메모리 부족 현상의 원인과 해결을 위한 아키텍처적 차이

일반적인 데이터 로딩 방식(In-memory)은 모든 샘플을 RAM에 한꺼번에 올립니다. 반면, 거대 데이터셋 전략은 참조(Reference)만 유지하고 실제 데이터는 연산 직전에만 메모리에 올리는 방식을 취합니다. 이 두 방식의 기술적 차이를 명확히 이해해야 병목을 해결할 수 있습니다.

2. 데이터 로딩 방식별 리소스 효율성 비교

성능과 메모리 사용량 사이의 트레이드오프를 분석한 표입니다.

구분	In-memory 로딩	Lazy Loading (Map-style)	Streaming (Iterable-style)
메모리 점유	매우 높음 (데이터 비례)	낮음 (배치 사이즈 비례)	최소 (단일 샘플 단위)
I/O 부하	초기 로드 시 집중	학습 내내 지속 발생	학습 내내 지속 발생
데이터 접근	무작위 접근(Random Access) 빠름	인덱싱 오버헤드 존재	순차 접근 위주
추천 데이터 크기	RAM 용량 이하	RAM 용량 초과 ~ 수백 GB	수 TB 이상의 분산 데이터
해결 과제	OOM 발생 원천 봉쇄 불가	SSD/I/O 병목 해결 필요	셔플링(Shuffling) 구현 복잡도

---

3. 실무 즉시 적용 가능한 대규모 데이터 로드 Example 7가지

실제 프로덕션 환경에서 테라바이트급 데이터를 안전하게 처리하기 위한 핵심 해결 코드입니다.

Example 1: 경로 기반의 Lazy Loading 구현 방법

이미지 데이터를 메모리에 올리지 않고 경로 리스트만 관리하여 __getitem__ 시점에 로드합니다.

import os
from PIL import Image
from torch.utils.data import Dataset

class LazyImageDataset(Dataset):
    def __init__(self, image_dir, transform=None):
        self.image_paths = [os.path.join(image_dir, f) for f in os.listdir(image_dir)]
        self.transform = transform

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

    def __getitem__(self, idx):
        # 호출 시점에만 디스크에서 읽어와 메모리 OOM 해결
        image = Image.open(self.image_paths[idx]).convert("RGB")
        if self.transform:
            image = self.transform(image)
        return image
        

Example 2: 메모리 맵(Memory Mapping)을 이용한 대형 NumPy/Tensor 로드

np.memmap을 사용하면 수십 GB의 바이너리 파일을 실제 RAM 사용 없이 배열처럼 다룰 수 있습니다.

import numpy as np
import torch

# 디스크 상의 거대 파일을 메모리에 매핑
large_data = np.memmap('large_features.dat', dtype='float32', mode='r', shape=(1000000, 512))

class MemMapDataset(Dataset):
    def __getitem__(self, idx):
        return torch.from_numpy(large_data[idx].copy())
        

Example 3: h5py 라이브러리를 이용한 계층적 데이터 로드

HDF5 포맷은 수억 개의 샘플을 압축 보관하며 효율적인 슬라이싱을 지원합니다.

import h5py

class HDF5Dataset(Dataset):
    def __init__(self, file_path):
        self.file = h5py.File(file_path, 'r')
        self.data = self.file['data']

    def __len__(self):
        return self.data.shape[0]

    def __getitem__(self, idx):
        return torch.from_numpy(self.data[idx])
        

Example 4: IterableDataset을 이용한 데이터 스트리밍 해결

네트워크 드라이브나 대형 클라우드 스토리지에서 데이터를 흐르듯이 읽어옵니다.

from torch.utils.data import IterableDataset

class StreamDataset(IterableDataset):
    def __init__(self, file_path):
        self.file_path = file_path

    def __iter__(self):
        with open(self.file_path, 'r') as f:
            for line in f:
                yield self.process_line(line)
        

Example 5: WebDataset을 이용한 샤드(Shard) 단위 로딩

데이터를 수천 개의 .tar 파일로 쪼개어 로컬 캐싱과 함께 병렬 로딩합니다.

# pip install webdataset
import webdataset as wds

dataset = wds.WebDataset("dataset_shard-{000..999}.tar").shuffle(1000).decode("rgb").to_tuple("jpg", "json")
loader = torch.utils.data.DataLoader(dataset, batch_size=64)
        

Example 6: LMDB(Lightning Memory-Mapped Database) 활용 방법

작은 이미지나 텍스트 데이터가 수백만 개일 때 파일 시스템 오버헤드를 줄이는 해결책입니다.

import lmdb

class LMDBDataset(Dataset):
    def __init__(self, db_path):
        self.env = lmdb.open(db_path, readonly=True, lock=False)
        
    def __getitem__(self, idx):
        with self.env.begin() as txn:
            byte_data = txn.get(str(idx).encode())
            # 바이트 데이터를 텐서로 역직렬화
        return data
        

Example 7: Parquet 포맷과 Pandas Chunking을 이용한 정형 데이터 처리

import pandas as pd

def load_large_csv(file_path):
    # 청크 단위로 읽어 메모리 점유율을 일정하게 유지
    for chunk in pd.read_csv(file_path, chunksize=10000):
        yield torch.tensor(chunk.values)
        

---

4. 대규모 데이터 처리 시 성능 병목 해결 팁

1. I/O 대역폭 확보: Lazy Loading 시에는 HDD보다 NVMe SSD 사용이 필수적입니다. I/O가 연산을 못 따라가면 GPU 대기 시간이 길어집니다.
2. Prefetching 활용: DataLoader의 prefetch_factor를 설정하여 CPU가 다음 배치를 미리 준비하도록 하십시오.
3. 공유 메모리 설정: Docker 환경에서는 --shm-size를 충분히 확보하여 멀티프로세싱 워커 간 데이터 전달 오류를 방지하십시오.

5. 결론: 규모의 경제를 지탱하는 엔지니어링의 힘

PyTorch에서 거대한 데이터셋을 다루는 것은 단순히 문법의 문제가 아니라 하드웨어 리소스와 알고리즘 사이의 균형을 맞추는 최적화 작업입니다. Lazy Loading과 IterableDataset을 적재적소에 배치하여 메모리 한계를 극복하고, h5py나 LMDB 같은 고성능 포맷을 도입함으로써 테라바이트 시대의 딥러닝 리더가 되시기 바랍니다.

내용 출처 및 기술 참조

• PyTorch Documentation: "Efficient Data Loading" & "IterableDataset API"
• "Training Deep Learning Models with Large Datasets" (Google Cloud Architecture Center)
• WebDataset Project GitHub: High-performance I/O for Large Scale Deep Learning