PyTorch Lightning 中的数据迭代器支持详解-优快云博客

本文链接：https://blog.youkuaiyun.com/gitblog_00506/article/details/148360245

PyTorch Lightning 中的数据迭代器支持详解

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

什么是数据迭代器

在深度学习中，数据迭代器(iterable)是指能够被循环遍历的数据对象。Python中的列表(list)、字典(dict)等都是常见的迭代器。在PyTorch生态中，torch.utils.data.DataLoader是最常用的数据迭代器，它通常从Dataset或IterableDataset中获取数据。

PyTorch Lightning框架的一个强大特性是它对任意迭代器的广泛支持，这使得开发者可以灵活地处理各种数据加载场景。

基本迭代器支持

PyTorch Lightning的Trainer类可以与任何Python迭代器协同工作，不过大多数情况下开发者还是会使用DataLoader作为主要的数据供给方式。这种设计提供了极大的灵活性，例如：

# 直接返回一个DataLoader
return DataLoader(...)

# 甚至可以直接返回一个简单的range迭代器
return list(range(1000))

多迭代器组合支持

在实际应用中，我们经常需要同时处理多个数据源。PyTorch Lightning对此提供了优雅的支持，允许以多种方式组合多个迭代器：

# 字典形式 - 生成批次格式为{'a': batch_a, 'b': batch_b}
return {"a": DataLoader(...), "b": DataLoader(...)}

# 列表形式 - 生成批次格式为[batch1, batch2]
return [DataLoader(...), DataLoader(...)]

# 更复杂的嵌套组合
return {"a": [dl1, dl2], "b": [dl3, dl4]}

批次组合模式

当使用多个迭代器时，PyTorch Lightning会自动根据"模式"来组合批次。这一功能由CombinedLoader类实现，提供了多种组合策略：

max_size_cycle：默认训练模式，循环遍历所有迭代器
min_size：以最短的迭代器为基准
sequential：默认验证/测试模式，顺序处理各迭代器

开发者可以显式指定组合模式：

from lightning.pytorch.utilities import CombinedLoader

iterables = {"a": DataLoader(), "b": DataLoader()}
combined_loader = CombinedLoader(iterables, mode="min_size")

在LightningDataModule中使用

LightningDataModule是组织数据加载逻辑的理想场所，它支持为不同阶段定义多个数据加载器：

class DataModule(LightningDataModule):
    def train_dataloader(self):
        return DataLoader(self.train_dataset)

    def val_dataloader(self):
        return [DataLoader(self.val_dataset_1), DataLoader(self.val_dataset_2)]

    def test_dataloader(self):
        return DataLoader(self.test_dataset)

    def predict_dataloader(self):
        return DataLoader(self.predict_dataset)