在PyTorch中,数据载入是训练深度学习模型的重要一环。
本文将介绍三种常用的数据载入方式:Dataset、DataLoader、以及自定义的数据加载器。
- 使用 Dataset 载入数据
方法:
from torch.utils.data import Datasetclass CustomDataset(Dataset):def __init__(self, ...):# 初始化数据集# ...def __len__(self):# 返回数据集的大小return len(self.data)def __getitem__(self, idx):# 根据索引返回样本和标签return self.data[idx], self.labels[idx]
使用示例:
custom_dataset = CustomDataset(...)
- 使用 DataLoader 加载数据集
方法:
from torch.utils.data import DataLoaderdata_loader = DataLoader(dataset,batch_size=32,shuffle=True,num_workers=4,drop_last=True,
)
参数说明:
dataset: 要加载的数据集对象。
batch_size: 每个批次的样本数量。
shuffle: 是否在每个 epoch 开始时打乱数据。
num_workers: 用于数据加载的进程数。
drop_last: 是否丢弃最后一个不完整的批次。
使用示例:
for data, labels in data_loader:# 训练模型的逻辑
- 使用自定义的 Data Loader
方法:
from torch.utils.data import DataLoader, IterableDatasetclass CustomDataLoader(IterableDataset):def __init__(self, ...):# 初始化数据加载器# ...def __iter__(self):# 返回一个迭代器return iter(self.data)
使用示例:
custom_loader = CustomDataLoader(...)
for data in custom_loader:# 自定义数据处理逻辑
注意事项:
Dataset 和 DataLoader 是 PyTorch 提供的数据载入工具,通常能满足大多数情况。
自定义数据加载器 (IterableDataset) 可以用于特殊情况下的数据加载需求。
总结:
使用 Dataset 创建数据集对象,实现 len 和 getitem 方法。
使用 DataLoader 加载数据集,设置参数如 batch_size、shuffle 等。
可选地,使用自定义的数据加载器 (IterableDataset) 处理特殊情况下的数据载入需求。