Pytorch 实现数据集自定义读取

以下是使用 PyTorch 实现数据集自定义读取的完整攻略，步骤分为五个主要部分，分别是：

1. 继承 Dataset 类并实现 len 和 getitem 函数
2. 定义数据集的标签和图像路径
3. 对数据集进行预处理
4. 加载数据集并创建 DataLoader
5. 使用 DataLoader 进行训练

首先，我们需要导入 PyTorch 和相关的库：

import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image

1. 继承 Dataset 类并实现 len 和 getitem 函数

我们需要创建一个自定义的类来实现这个数据集，这个类需要继承 PyTorch 的 Dataset 类，并且实现 len 和 getitem 函数。len 函数需要返回数据集的大小，getitem 函数需要返回指定索引位置的图像和标签。

class CustomDataset(Dataset):
    def __init__(self, img_dir, labels, transform=None):
        self.img_dir = img_dir
        self.labels = labels
        self.transform = transform

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

    def __getitem__(self, idx):
        img_path = os.path.join(self.img_dir, f"{idx}.jpg")
        img = Image.open(img_path)

        if self.transform:
            img = self.transform(img)

        label = self.labels[idx]

        return img, label

在这个类中，我们需要传入四个参数，分别是：

• img_dir：图像文件夹路径
• labels：数据集标签
• transform：预处理的方式

在 len 函数中，我们直接返回了标签的长度。

在 getitem 函数中，我们首先通过索引将图像路径读取并打开，然后根据索引获取标签。最后，如果定义了 transform 函数，则对图像进行预处理。

2. 定义数据集的标签和图像路径

我们需要定义数据集的标签和图像路径，这些标签和路径将用于创建自定义数据集。

img_dir = "images/"
labels = [0, 1, 0, 1, 0, 1, 1, 0, 0, 1]

上面的代码片段中，我们设置了一个图像文件夹路径和一个标签列表。这里我们定义了 10 张图像和它们的标签。

3. 对数据集进行预处理

在定义自定义 Dataset 类时，我们传入了一个参数 transform，这个参数用于对数据集进行预处理。我们可以使用 PyTorch 提供的 transforms 库对图像进行常用的数据预处理操作。

下面是一个示例：

transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

这个 transform 对象将对图像进行三个常用的操作：图像大小调整，转换为张量，以及归一化。

4. 加载数据集并创建 DataLoader

通过上面的步骤，我们已经定义好了自定义数据集及其预处理方式。现在我们需要将数据集加载到 DataLoader 中，以便在训练过程中进行批量读取和处理。

dataset = CustomDataset(img_dir=img_dir, labels=labels, transform=transform)
dataloader = DataLoader(dataset=dataset, batch_size=2, shuffle=True)

在这个示例中，我们首先创建了一个 CustomDataset 对象，并传入了上面定义的图像文件夹路径、标签列表以及预处理 transform。然后我们使用 DataLoader 对象将数据集加载进来，设置了 batch_size 为 2，也就是每次读取两张图像，shuffle 为 True，用于打乱数据集顺序。

5. 使用 DataLoader 进行训练

现在我们已经准备好了自定义数据集和 DataLoader，在训练模型时，我们只需要循环迭代 DataLoader，并传入图像和标签即可。

以下是一个简单的训练示例：

for batch_idx, (data, target) in enumerate(dataloader):
    # 训练代码
    # ...

在每次循环迭代中，我们可以访问到一个批次的图像和标签。其中 data 和 target 分别代表图像和标签。

至此，我们就完成了使用 PyTorch 实现数据集自定义读取的攻略。