pytorch自制数据集载入

在 PyTorch 中，可以使用 `torch.utils.data.Dataset` 类来自定义数据集，并使用 `torch.utils.data.DataLoader` 类来加载数据集。

下面是一个简单的示例，假设我们有一个包含图像和标签的数据集，图像存储在 `images` 文件夹中，标签存储在 `labels.txt` 文件中，每行包含一个标签值（整数）：

import torch
from PIL import Image

class CustomDataset(torch.utils.data.Dataset):
    def __init__(self, root_dir, transform=None):
        self.root_dir = root_dir
        self.transform = transform
        self.labels = []
        with open('labels.txt', 'r') as f:
            lines = f.readlines()
            self.labels = [int(line.strip()) for line in lines]

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

    def __getitem__(self, idx):
        img_path = f'{self.root_dir}/{idx}.jpg'
        img = Image.open(img_path).convert('RGB')
        label = self.labels[idx]
        if self.transform:
            img = self.transform(img)
        return img, label

# 创建数据集实例
dataset = CustomDataset('images')
# 创建数据集加载器
dataloader = torch.utils.data.DataLoader(dataset, batch_size=32, shuffle=True)

在上面的示例中，`CustomDataset` 类定义了自定义数据集的加载方式。`__init__` 方法用于初始化数据集实例，读取标签文件并保存标签。`__len__` 方法返回数据集的大小，即样本数量。`__getitem__` 方法用于获取一个样本，其中 `idx` 表示样本的索引，根据索引读取该样本的图像文件和标签文件，并返回图像和标签。

我们使用 `dataset = CustomDataset('images')` 创建数据集实例，然后使用 `dataloader = torch.utils.data.DataLoader(dataset, batch_size=32, shuffle=True)` 创建数据集加载器。其中，`batch_size` 表示每个批次的样本数量，`shuffle` 表示是否随机打乱数据集。

使用 `dataloader` 加载数据集时，会自动调用 `CustomDataset` 中的 `__getitem__` 方法获取样本，然后将多个样本拼接成一个批次返回。因此，可以直接使用 `for batch_x, batch_y in dataloader:` 的方式遍历数据集。