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)

train DataLoader の作成

学習データのバッチを提供
GPU 使用率の最適化に役立つ
大規模データセットの効率的な反復処理を可能にする

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

validation DataLoader の作成

モデル検証用のデータを提供
汎化性能のモニタリングに役立つ
シャッフルにより評価の一貫性を確保

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

test DataLoader の作成

学習完了後の最終的なモデル評価用データを提供
実環境でのパフォーマンス評価をシミュレート
偏りのない性能測定を保証

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

DataModule と LightningModule を接続する

モジュール設計により、データとモデルのロジックを分離

PyTorch Lightning の図

DataModule と LightningModule を接続する

モジュール設計により、データとモデルのロジックを分離
LightningDataModule は LightningModule と連携する
標準化されたワークフローにより再現性が向上

DataModule と LightningModule を示す PyTorch Lightning の図

練習しましょう！

PyTorch Lightning によるスケーラブルな AI モデルの構築