PyTorch Lightning 中的 LightningDataModule 详解-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00571/article/details/148360241

PyTorch Lightning 中的 LightningDataModule 详解

pytorch-lightning Lightning-AI/pytorch-lightning: PyTorch Lightning 是一个轻量级的高级接口，用于简化 PyTorch 中深度学习模型的训练流程。它抽象出了繁杂的工程细节，使研究者能够专注于模型本身的逻辑和实验设计，同时仍能充分利用PyTorch底层的灵活性。项目地址: https://gitcode.com/gh_mirrors/py/pytorch-lightning

什么是 LightningDataModule

LightningDataModule 是 PyTorch Lightning 框架中用于管理数据的一个核心组件。它将数据处理流程标准化、模块化，使得数据准备工作更加清晰和可复用。

简单来说，LightningDataModule 封装了 PyTorch 数据处理中的五个关键步骤：

数据下载/预处理
数据清洗与保存
数据集加载
数据转换/增强
数据加载器包装

为什么需要 DataModule

在传统的 PyTorch 项目中，数据处理代码通常分散在多个文件中，这会导致以下问题：

难以在不同项目间复用相同的数据处理流程
难以确保不同实验使用相同的数据划分和转换
数据预处理细节（如标准化参数）难以追踪
团队协作时难以共享完整的数据处理方案

LightningDataModule 通过将所有这些逻辑封装在一个类中，完美解决了这些问题。

DataModule 的核心结构

一个完整的 LightningDataModule 通常包含以下方法：

prepare_data()

负责一次性执行的操作，如：

数据下载
数据预处理
数据保存到磁盘

这个方法只在主进程上执行一次，确保在多进程环境下不会重复下载或处理数据。

setup(stage: str)

负责每个进程需要执行的操作，如：

数据划分（训练/验证/测试集）
数据集创建
数据转换应用
特征提取

stage 参数用于区分不同的训练阶段（"fit"、"test"、"predict"等）。

数据加载器方法

train_dataloader(): 返回训练数据加载器
val_dataloader(): 返回验证数据加载器
test_dataloader(): 返回测试数据加载器
predict_dataloader(): 返回预测数据加载器

teardown(stage: str)

用于清理资源，在训练/测试/预测结束后调用。

实际应用示例

下面是一个完整的 MNIST 数据模块实现：

import lightning as L
from torch.utils.data import DataLoader, random_split
from torchvision.datasets import MNIST
from torchvision import transforms

class MNISTDataModule(L.LightningDataModule):
    def __init__(self, data_dir="./", batch_size=32):
        super().__init__()
        self.data_dir = data_dir
        self.batch_size = batch_size
        self.transform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize((0.1307,), (0.3081,))
        ])

    def prepare_data(self):
        # 下载数据（只执行一次）
        MNIST(self.data_dir, train=True, download=True)
        MNIST(self.data_dir, train=False, download=True)

    def setup(self, stage=None):
        # 数据划分和转换
        if stage == "fit" or stage is None:
            mnist_full = MNIST(self.data_dir, train=True, transform=self.transform)
            self.mnist_train, self.mnist_val = random_split(
                mnist_full, [55000, 5000], 
                generator=torch.Generator().manual_seed(42)
            )
        
        if stage == "test" or stage is None:
            self.mnist_test = MNIST(self.data_dir, train=False, transform=self.transform)

    def train_dataloader(self):
        return DataLoader(self.mnist_train, batch_size=self.batch_size)

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

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

如何使用 DataModule

使用 DataModule 非常简单：

# 初始化
dm = MNISTDataModule()
model = MyLightningModel()

# 训练
trainer = L.Trainer()
trainer.fit(model, datamodule=dm)

# 测试
trainer.test(datamodule=dm)

高级功能

超参数保存

DataModule 也支持超参数保存功能：

class CustomDataModule(L.LightningDataModule):
    def __init__(self, batch_size=32, learning_rate=1e-3):
        super().__init__()
        self.save_hyperparameters()
        
    def setup(self, stage):
        # 可以使用 self.hparams 访问保存的参数
        print(f"Batch size: {self.hparams.batch_size}")

自定义数据转换

DataModule 可以灵活处理各种数据转换需求：

def __init__(self):
    self.train_transform = transforms.Compose([
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor()
    ])
    self.val_transform = transforms.ToTensor()