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

本文链接：https://blog.youkuaiyun.com/gitblog_00023/article/details/148360182

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

pytorch-lightning 项目地址: https://gitcode.com/gh_mirrors/pyt/pytorch-lightning

概述

在深度学习项目中，数据加载和处理是模型训练的关键环节。PyTorch Lightning作为一个高级PyTorch封装框架，提供了对Python可迭代对象的全面支持，使得数据加载更加灵活高效。本文将深入探讨PyTorch Lightning中可迭代数据集的使用方法。

什么是可迭代对象

在Python中，可迭代对象(iterable)是指实现了__iter__()方法的对象，能够被循环遍历。常见的可迭代对象包括：

列表(list)
字典(dict)
集合(set)
字符串(str)
生成器(generator)

在PyTorch生态中，torch.utils.data.DataLoader是最常用的可迭代对象，它通常从Dataset或IterableDataset中获取数据。

PyTorch Lightning对可迭代对象的支持

PyTorch Lightning的Trainer类能够与任意可迭代对象协同工作，不过大多数用户还是会选择使用DataLoader作为主要的数据加载方式。

多数据加载器支持

PyTorch Lightning的一个强大特性是支持多个数据加载器的组合使用。以下是几种常见的组合方式：

# 单个DataLoader
return DataLoader(...)

# 简单可迭代对象
return list(range(1000))

# 字典形式的DataLoader
# 产生的batch格式：{'a': loader_a的batch, 'b': loader_b的batch}
return {"a": DataLoader(...), "b": DataLoader(...)}

# 列表形式的DataLoader
# 产生的batch格式：[loader_1的batch, loader_2的batch]
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")
trainer.fit(model, combined_loader)

注意：当前trainer.predict方法仅支持sequential模式，而trainer.fit方法不支持该模式。

使用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)

使用LightningModule钩子

同样的数据加载器配置也可以在LightningModule中实现：

class MyModel(LightningModule):
    def train_dataloader(self):
        return DataLoader(...)
    
    # 其他数据加载器方法...

直接向Trainer传递数据加载器

PyTorch Lightning的Trainer方法也支持直接传递可迭代对象或可迭代对象集合：

trainer.fit(model, train_dataloader, val_dataloaders)
trainer.validate(model, val_dataloaders)
trainer.test(model, test_dataloaders)
trainer.predict(model, predict_dataloaders)