Pytorch建模过程中的DataLoader与Dataset示例详解

PyTorch是一个非常流行的深度学习框架， 绝大多数项目中都需要使用数据加载器(DataLoader)来加载模型训练所需的数据。在这篇攻略中，我们将详细讲解如何使用数据集(Dataset)和数据加载器(DataLoader)来加载训练数据。

什么是数据集(Dataset)？

在PyTorch中，数据集被定义为一个抽象类(torch.utils.data.Dataset)，我们需要继承它并根据我们自己的数据集来实现它。数据集必须实现两个方法： __len__和__getitem__。

__len__方法

__len__方法返回数据集中样本数量。例如，如果您的数据集有100张图片，则__len__应该返回100。

__getitem__方法

__getitem__方法负责将索引转换为数据集中的样本。通常，它从磁盘中加载数据并返回一个tensor 。例如，如果您有一个包含图像和相应标签的数据集，则__getitem__方法应该返回图像和对应标签。

示例1

让我们以一个简单的例子开始，假设我们有一个CSV格式的数据文件，其中包含每个样本的图像路径和相应标签。我们需要读取CSV文件，并从磁盘中读取图像和标签。我们来看一下如何为此实现一个自定义数据集(Dataset)

首先是CSV数据文件的格式

path,label
data/0001.png,1
data/0002.png,0
data/0003.png,0
.....

下面是我们实现的例子：

import pandas as pd
import torch
from PIL import Image
from torch.utils.data import Dataset

class CustomDataset(Dataset):
    def __init__(self, csv_file):
        self.data = pd.read_csv(csv_file)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        path = self.data.iloc[index, 0]
        label = self.data.iloc[index, 1]
        img = Image.open(path)
        img = img.convert("RGB")
        img_tensor = torch.tensor(img)
        return img_tensor, label

在上面的代码中，我们首先用 pandas 读取CSV数据文件。 然后，在__len__方法中，我们返回数据集中的总样本数。最后，在__getitem__方法中，我们从数据集中读取一张图片，并将其转换为torch.tensor。获取的图像tensor以及它的标签是作为一个元组返回的。

什么是数据加载器(DataLoader)?

在上面的示例中，我们已经实现了一个自定义数据集(Dataset)，但以这种方式读取数据并不是我们需要的。 还需要将数据加载进模型中进行训练。我们需要使用数据加载器(DataLoader)

数据加载器(DataLoader)是PyTorch中的一个迭代器，可以对任意数据集进行批量处理、并行加载和数据重组。在对模型进行训练之前，数据集被加载到数据加载器中。数据集在每个纪元(epoch)中都会被重新加载，并且数据加载器将为每个批处理提供数据。

数据加载器(DataLoader)具有以下常用参数：

• dataset: 用于加载数据的数据集对象。
• batch_size: 批量大小。
• shuffle: 是否要对数据进行随机重组。
• num_workers: 使用的子进程数量。

示例2

现在，我们已经实现了自定义数据集(Dataset)，接下来，我们将通过数据加载器(DataLoader)来加载数据并对其进行处理

from torch.utils.data import DataLoader

dataset = CustomDataset("data_file.csv")
dataloader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4)

在上面的代码中，我们首先指定CustomDataset作为我们的数据集，然后使用DataLoader来加载数据集，并设置批量大小为32，随机重组数据并使用4个子进程来加载数据。

总结

在本文中，我们已经详细讲解了PyTorch中数据集(Dataset)和数据加载器(DataLoader)的用法。实现自定义数据集并初始化数据加载器可以帮助您快速、高效地加载训练数据。在训练模型的过程中，数据集和数据加载器是非常重要的组成部分，这些技巧将有助于您快速地开始使用PyTorch进行模型训练。