当前位置：首页 > news >正文

DataLoader 的基本用法

news 来源：原创 2024/9/20 2:46:33

DataLoader

在深度学习中，数据加载和预处理是训练模型的关键步骤。PyTorch 提供了 DataLoader 用于简化这一过程。本文将详细介绍 PyTorch 中 DataLoader 的使用，包括基本用法、常见参数及自定义数据集的方式。

基本概念

什么是 DataLoader

DataLoader 是 PyTorch 中的一个类，用于将数据集（通常是一个 Dataset 对象）打包成一个可迭代的对象，方便在训练过程中逐批次读取数据。DataLoader 可以处理数据的随机打乱、并行加载、多线程加载等。

DataLoader 的基本构造

DataLoader 从数据集中取样本，并能够在多个线程中异步读取数据，这是其设计的关键点。


from torch.utils.data import DataLoaderdataloader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4)

dataset: 一个继承自 Dataset 的对象，定义了如何获取数据。
batch_size: 指定了每个批次的数据量。
shuffle: 是否在每轮迭代时打乱数据。
num_workers: 加载数据时使用的子进程数，默认值为0，即使用主进程。

使用 DataLoader 加载数据

示例数据集
使用一个简单的数据集来演示 DataLoader 的基本用法。


from torch.utils.data import DataLoader, Datasetimport torchclass SimpleDataset(Dataset):def __init__(self):self.data = torch.arange(100).float().unsqueeze(1)  # 100个样本每个样本包含一个特征self.labels = torch.arange(100).float()  # 标签与数据相同def __len__(self):return len(self.data)def __getitem__(self, idx):x = self.data[idx]y = self.labels[idx]return x, ydataset = SimpleDataset()

创建 DataLoader


dataloader = DataLoader(dataset, batch_size=10, shuffle=True, num_workers=2)for batch in dataloader:inputs, labels = batchprint(inputs, labels)

DataLoader 的常见参数

batch_size
批次大小，默认为1。指每次迭代所返回的数据量。


dataloader = DataLoader(dataset, batch_size=32)

shuffle
是否在每个 Epoch 开始时打乱数据。默认为 False。


dataloader = DataLoader(dataset, batch_size=32, shuffle=True)

num_workers
加载数据时使用的子进程数。默认为 0，即在主进程中加载数据。如果设置为一个大于 0 的数值，则会使用多个进程来加载数据。


dataloader = DataLoader(dataset, batch_size=32, num_workers=4)

自定义 Dataset

有时，内置的数据集无法满足我们的需求，此时我们需要自定义数据集。自定义 Dataset 需要继承 torch.utils.data.Dataset 类并重写 len 和 getitem 方法。


class CustomDataset(Dataset):def __init__(self, data, labels):self.data = dataself.labels = labelsdef __len__(self):return len(self.data)def __getitem__(self, idx):x = self.data[idx]y = self.labels[idx]return x, ydata = torch.randn(100, 3)  # 100个样本，每个样本包含3个特征labels = torch.randint(0, 2, (100,))  # 100个标签，值为0或1custom_dataset = CustomDataset(data, labels)dataloader = DataLoader(custom_dataset, batch_size=10, shuffle=True)for batch in dataloader:inputs, labels = batchprint(inputs, labels)

数据转换

在实际应用中，通常需要对数据进行一定的预处理，这时可以使用 torchvision.transforms


from torchvision import transformsclass TransformedDataset(Dataset):def __init__(self, data, labels, transform=None):self.data = dataself.labels = labelsself.transform = transformdef __len__(self):return len(self.data)def __getitem__(self, idx):x = self.data[idx]y = self.labels[idx]if self.transform:x = self.transform(x)return x, ytransform = transforms.Compose([transforms.Normalize(mean=[0.5], std=[0.5])])data = torch.randn(100, 1)labels = torch.randint(0, 2, (100,))transformed_dataset = TransformedDataset(data, labels, transform=transform)dataloader = DataLoader(transformed_dataset, batch_size=10, shuffle=True)for batch in dataloader:inputs, labels = batchprint(inputs, labels)