LightningDataModule로 데이터 관리

PyTorch Lightning으로 만드는 확장 가능한 AI 모델

Sergiy Tkachuk

Director, GenAI Productivity

모델 학습을 위한 데이터 준비

데이터 준비가 미흡하면 학습 문제가 발생
- 학습이 느림
- 잦은 중단
- 수렴 실패

학습을 위한 데이터 준비.png

왜 LightningDataModule을 쓰나요?

📂 데이터셋 처리를 중앙화

📊 데이터 준비 워크플로 표준화

🚀 학습·평가 단계 단순화

LightningDataModule로 데이터 관리

핵심 메서드:

prepare_data: 데이터 다운로드 및 설정
setup: 학습/검증/테스트 분할

class ImageDataModule(pl.LightningDataModule):
    def __init__(self, data_dir="./data", batch_size=32):
        super().__init__()
        ...

    def prepare_data(self):
        datasets.MNIST(self.data_dir, train=True, download=True)

    def setup(self, stage=None):
        dataset = datasets.MNIST(self.data_dir, train=True, transform=self.transform)
        self.train_data, self.val_data = random_split(dataset, [55000, 5000])
        self.test_data = datasets.MNIST(self.data_dir, train=False, transform=self.transform)

학습 DataLoader 생성

학습 데이터 배치 제공
GPU 활용 최적화에 기여
대용량 데이터셋 효율적 반복 처리

def train_dataloader(self):
    return DataLoader(self.train_data, batch_size=self.batch_size, shuffle=True)

검증 DataLoader 생성

모델 검증용 데이터를 제공
일반화 성능 모니터링 지원
셔플링으로 평가 일관성 보장

def val_dataloader(self):
    return DataLoader(self.val_data, batch_size=self.batch_size)

테스트 DataLoader 생성

학습 완료 후 최종 평가 데이터 제공
실제 환경의 성능 평가를 모사
편향 없는 성능 측정 보장

def test_dataloader(self):
    return DataLoader(self.test_data, batch_size=self.batch_size)

DataModule을 LightningModule에 연결

모듈식 설계로 데이터와 모델 로직을 분리

PyTorch Lightning diagram

DataModule을 LightningModule에 연결

모듈식 설계로 데이터와 모델 로직 분리
LightningDataModule은 LightningModule과 짝을 이룸
표준화된 워크플로로 재현성 향상

DataModule과 LightningModule이 포함된 PyTorch Lightning 다이어그램

연습해 봅시다!

PyTorch Lightning으로 만드는 확장 가능한 AI 모델