PyTorch Lightning LightningDataModule 介绍

LightningDataModule 是 PyTorch Lightning 提供的数据模块，用于统一管理数据加载流程（包括数据准备、预处理、拆分、批量加载等）。它的核心作用是将数据处理逻辑与模型解耦，提高代码的可复用性和可读性。

---

1. LightningDataModule 的作用

✅ 封装数据预处理：数据下载、清理、转换等步骤都可以在 LightningDataModule 中完成。
✅ 统一数据加载流程：确保训练、验证、测试和推理数据集使用相同的数据预处理逻辑。
✅ 简化 Trainer 代码：LightningDataModule 使 Trainer.fit() 更加简洁和模块化。
✅ 支持多 GPU、TPU 训练：可以轻松适配不同计算设备的 Dataloader 设定。

---

2. LightningDataModule 的基本结构

LightningDataModule 主要包含以下关键方法：

方法	作用
prepare_data()	仅在主进程中运行一次，用于下载数据、处理静态数据（如数据去重）
setup(stage)	在每个 GPU/TPU 设备上运行，用于数据拆分（