---

PyTorch 自定义数据集示例

两年后的更新： 我创建这个仓库来指导刚开始接触 PyTorch 的人（就像当年的我一样）已经过去很久了。然而，随着岁月的流逝和不同项目的开展，我构建数据集的方式也多次改变。我在 这里 添加了一个更基础的数据集，以展示我现在常用的实现方式。我想强调的是，自定义数据集之所以被称为“自定义”，正是因为它可以根据你的需求进行任意设计。因此，在你参与不同的项目后，自然而然会形成自己独特的数据集构建方式。本项目中提供的示例并不是创建自定义数据集的唯一正确方法，而是为了展示 PyTorch 数据集的灵活性和可能性。希望这个仓库能够帮助你更好地理解 PyTorch 数据集。

PyTorch 官方仓库中也有一些自定义数据集的示例，比如 这个，但对于初学者来说（像当时的我一样），这些示例仍然有些晦涩难懂。我不得不花一些时间去弄清楚究竟需要做些什么才能构建一个完全自定义的数据集。为了避免大家翻阅大量资料的麻烦，我在这里总结了 PyTorch 数据集的基础知识。具体内容如下：

主题

• 自定义数据集基础
• 使用 torchvision transforms
• 另一种使用 torchvision transforms 的方式
• 集成 pandas（读取 CSV 文件）
• 在 __getitem__() 中加入更多逻辑的 pandas 集成
• 将类别信息嵌入文件名
• 是否使用 torchvision transforms？
• 结合自定义数据集使用数据加载器
• 未来的更新（希望如此）

自定义数据集基础

首先也是最重要的一点，就是创建一个数据集类。

from torch.utils.data.dataset import Dataset

class MyCustomDataset(Dataset):
    def __init__(self, ...):
        # 初始化操作
        
    def __getitem__(self, index):
        # 处理数据
        return (img, label)

    def __len__(self):
        return count # 你拥有的样本数量

这就是你需要填充的基本框架，才能构建一个自定义数据集。数据集必须包含以下方法，以便后续被数据加载器使用：

• __init__() 方法用于执行初始化逻辑，例如读取 CSV 文件、设置变换、筛选数据等。
• __getitem__() 方法返回数据及其标签。该方法会在数据加载器中被调用，例如：

img, label = MyCustomDataset.__getitem__(99)  # 获取第 99 个样本

这里的 MyCustomDataset 返回了一张图像和一个标签，但这并不意味着 __getitem__() 只能返回这两样内容。根据你的应用场景，你可以返回多种类型的数据。

需要注意的是，__getitem__() 必须为单个数据点返回特定类型的对象（如张量、NumPy 数组等），否则在数据加载器中会遇到类似以下的错误：

TypeError: batch must contain tensors, numbers, dicts or lists; found <class 'PIL.PngImagePlugin.PngImageFile'>

• __len__() 方法返回数据集中样本的总数。

使用 torchvision transforms

在大多数示例中，__init__() 方法中的 transforms 参数通常被设置为 None。这是用来对你的数据或图像应用 torchvision 提供的变换操作。详细的变换列表可以参考 这里 和 这里。最常见的用法如下：

from torch.utils.data.dataset import Dataset
from torchvision import transforms

class MyCustomDataset(Dataset):
    def __init__(self, ..., transforms=None):
        # 初始化操作
        ...
        self.transforms = transforms
        
    def __getitem__(self, index):
        # 处理数据
        ...
        data = # 从文件或图像中读取的数据
        if self.transforms is not None:
            data = self.transforms(data)
        # 如果 transforms 不为空，则按照其创建顺序依次应用变换操作
        return (img, label)

    def __len__(self):
        return count # 你拥有的数据（图像）数量
        
if __name__ == '__main__':
    # 定义变换 (1)
    transformations = transforms.Compose([transforms.CenterCrop(100), transforms.ToTensor()])
    # 实例化数据集
    custom_dataset = MyCustomDataset(..., transformations)
    

另一种使用 torchvision transforms 的方式

就我个人而言，我不太喜欢把数据集的变换逻辑放在数据集类之外（见上面的 (1)）。因此，除了上述示例中的用法外，你还可以这样使用：

from torch.utils.data.dataset import Dataset
from torchvision import transforms

class MyCustomDataset(Dataset):
    def __init__(self, ...):
        # 初始化操作
        ...
        # (2) 一种方法是单独定义每个变换
        # 当你定义变换时，会调用该变换的 __init__() 方法
        self.center_crop = transforms.CenterCrop(100)
        self.to_tensor = transforms.ToTensor()
        
        # (3) 或者你也可以直接组合它们
        self.transformations = \
            transforms.Compose([transforms.CenterCrop(100),
                                transforms.ToTensor()])
        
    def __getitem__(self, index):
        # 处理数据
        ...
        data = # 从文件或图像中读取的数据
        
        # 第二次调用变换时，会触发 __call__() 方法并应用变换
        data = self.center_crop(data)  # (2)
        data = self.to_tensor(data)  # (2)
        
        # 或者你可以直接使用组合好的变换
        data = self.transformations(data)  # (3)
        
        # 注意，你只需要选择其中一种实现方式即可，(2) 或 (3)
        return (img, label)

    def __len__(self):
        return count # 你拥有的数据（图像）数量
        
if __name__ == '__main__':
    # 实例化数据集
    custom_dataset = MyCustomDataset(...)
    

结合 Pandas

假设我们想用 Pandas 从一个 CSV 文件中读取数据。第一个示例的 CSV 文件如下（虽然没有表头，但这并不影响），它包含文件名、标签（类别）以及一个额外的操作指示符。根据这个操作标志，我们会对图像执行一些操作。

文件名	标签	额外操作
tr_0.png	5	TRUE
tr_1.png	0	FALSE
tr_1.png	4	FALSE

如果我们想要构建一个自定义数据集，从这个 CSV 文件中读取图像路径，可以按照以下方式实现：

class CustomDatasetFromImages(Dataset):
    def __init__(self, csv_path):
        """
        参数:
            csv_path (string): CSV 文件的路径
            img_path (string): 图像所在文件夹的路径
            transform: PyTorch 的变换工具，用于图像变换和转换为张量
        """
        # 定义转换：将图像转换为张量
        self.to_tensor = transforms.ToTensor()
        # 读取 CSV 文件
        self.data_info = pd.read_csv(csv_path, header=None)
        # 第一列是图像路径
        self.image_arr = np.asarray(self.data_info.iloc[:, 0])
        # 第二列是标签
        self.label_arr = np.asarray(self.data_info.iloc[:, 1])
        # 第三列是操作指示符
        self.operation_arr = np.asarray(self.data_info.iloc[:, 2])
        # 计算数据长度
        self.data_len = len(self.data_info.index)

    def __getitem__(self, index):
        # 从 Pandas DataFrame 中获取图像名称
        single_image_name = self.image_arr[index]
        # 打开图像
        img_as_img = Image.open(single_image_name)

        # 检查是否有操作
        some_operation = self.operation_arr[index]
        # 如果有操作
        if some_operation:
            # 对图像执行某些操作
            # ...
            # ...
            pass
        # 将图像转换为张量
        img_as_tensor = self.to_tensor(img_as_img)

        # 根据截取的 Pandas 列获取图像的标签（类别）
        single_image_label = self.label_arr[index]

        return (img_as_tensor, single_image_label)

    def __len__(self):
        return self.data_len

if __name__ == "__main__":
    # 调用数据集
    custom_mnist_from_images =  \
        CustomDatasetFromImages('../data/mnist_labels.csv')

结合 Pandas 并加入更多逻辑

另一个例子可能是从 CSV 文件中读取图像，其中每个像素的值都列在单独的一列中。（有时 MNIST 数据集就是以这种方式提供的）。这只需要修改 __getitem__() 方法中的逻辑即可。最终，你只需返回图像的张量及其标签。数据可能被组织成如下形式：

标签	pixel_1	pixel_2	...
1	50	99	...
0	21	223	...
9	44	112	...

class CustomDatasetFromCSV(Dataset):
    def __init__(self, csv_path, height, width, transforms=None):
        """
        参数:
            csv_path (string): CSV 文件的路径
            height (int): 图像的高度
            width (int): 图像的宽度
            transform: PyTorch 的变换工具，用于图像变换和转换为张量
        """
        self.data = pd.read_csv(csv_path)
        self.labels = np.asarray(self.data.iloc[:, 0])
        self.height = height
        self.width = width
        self.transforms = transform

    def __getitem__(self, index):
        single_image_label = self.labels[index]
        # 读取每行的 784 个像素值，并将一维数组 ([784]) 重塑为二维数组 ([28,28])
        img_as_np = np.asarray(self.data.iloc[index][1:]).reshape(28,28).astype('uint8')
	# 将 NumPy 数组转换为 PIL 图像，模式 'L' 表示灰度图
        img_as_img = Image.fromarray(img_as_np)
        img_as_img = img_as_img.convert('L')
        # 将图像转换为张量
        if self.transforms is not None:
            img_as_tensor = self.transforms(img_as_img)
        # 返回图像和标签
        return (img_as_tensor, single_image_label)

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

if __name__ == "__main__":
    transformations = transforms.Compose([transforms.ToTensor()])
    custom_mnist_from_csv = \
        CustomDatasetFromCSV('../data/mnist_in_csv.csv', 28, 28, transformations)
        

一个自定义的自定义的自定义数据集

我已经很久没有更新这个仓库了（居然有两年了），在这段时间里，我几乎完全不再使用 torchvision 的变换以及 CSV 等文件格式（除非绝对必要）。我之所以不再使用 torchvision 的变换，一方面是因为我自己编写了相应的变换函数；但更重要的原因是，我不太喜欢在大多数最佳实践示例中，将变换作为参数传入数据集类并在初始化时设置的做法，因为这并不是一种最佳的方式。在我的训练数据集中，我经常需要对图像或数据进行一些随机的预处理操作（如翻转、镜像、填充、添加噪声、饱和度调整、裁剪等），并且希望能够自由地选择这些操作的程度，甚至决定是否应用它们。因此，我的数据集通常遵循如下流程：

1- 获取待读取图像/数据的位置
2- 读取图像/数据
3- 转换为 NumPy 数组
4- 对 NumPy 数组进行一些处理（随机）
5- 对 NumPy 数组进行一些处理（随机）
6- 对 NumPy 数组进行一些处理（随机）
...
15- 将数据转换为张量
返回数据位置、数据名称、数据和标签

当然，你也可以像我上面列出的那样直接应用变换，归根结底这只是个人偏好问题。

下面我分享一个我在大多数实验中使用的极简自定义数据集。

class CustomDatasetFromFile(Dataset):
    def __init__(self, folder_path):
        """
        一个数据集示例，其中类别信息嵌入在文件名中。
        此数据集也不使用任何 PyTorch 变换。

        Args:
            folder_path (string): 图像文件夹路径
        """
        # 获取图像列表
        self.image_list = glob.glob(folder_path+'*')
        # 计算数据长度
        self.data_len = len(self.image_list)

    def __getitem__(self, index):
        # 从文件路径中获取图像名称
        single_image_path = self.image_list[index]
        # 打开图像
        im_as_im = Image.open(single_image_path)

        # 对图像进行一些操作
        # 转换为 NumPy 数组，尺寸为 28x28
        im_as_np = np.asarray(im_as_im)/255
        # 增加通道维度，尺寸为 1x28x28
        # 注意：如果你读取的是 RGB 图像，或者数据本身已经有通道维度，则无需执行此步骤。
        im_as_np = np.expand_dims(im_as_np, 0)
        # 对 NumPy 数组进行一些预处理操作
        # ...
        # ...
        # ...

        # 将图像转换为张量，并更改数据类型
        im_as_ten = torch.from_numpy(im_as_np).float()

        # 根据文件名获取图像的标签（类别）
        class_indicator_location = single_image_path.rfind('_c')
        label = int(single_image_path[class_indicator_location+2:class_indicator_location+3])
        return (im_as_ten, label)

    def __len__(self):
        return self.data_len

使用数据加载器

PyTorch 的数据加载器会调用 __getitem__() 方法，并将数据打包成批次。理论上，我们也可以不使用数据加载器，而是逐个调用 __getitem__() 方法并将数据输入模型（尽管使用数据加载器确实非常方便）。继续以上面的例子为例，假设我们有一个名为 CustomDatasetFromCSV 的自定义数据集，那么我们可以这样调用数据加载器：

...
if __name__ == "__main__":
    # 定义变换
    transformations = transforms.Compose([transforms.ToTensor()])
    # 定义自定义数据集
    custom_mnist_from_csv = \
        CustomDatasetFromCSV('../data/mnist_in_csv.csv',
                             28, 28,
                             transformations)
    # 定义数据加载器
    mn_dataset_loader = torch.utils.data.DataLoader(dataset=custom_mnist_from_csv,
                                                    batch_size=10,
                                                    shuffle=False)
    
    for images, labels in mn_dataset_loader:
        # 将数据输入模型

数据加载器的第一个参数是数据集对象，它会调用该数据集的 __getitem__() 方法。batch_size 参数决定了每个批次包含多少个独立的数据点。如果假设单个图像张量的尺寸为 1x28x28（D:1, H:28, W:28），那么通过这个数据加载器返回的张量将是 10x1x28x28（批次-深度-高度-宽度）。

关于 多 GPU 的使用说明。在 PyTorch 中，多 GPU 数据加载器的工作方式是尝试将批次均匀地分配到所有可用的 GPU 上。因此，如果使用的 batch_size 小于 GPU 的数量，就无法充分利用所有的 GPU。

未来更新

我会在有空的时候继续更新这个仓库。以下是我计划加入的一些内容。如果你希望看到其他特定的示例，请随时告诉我。

• 一个可运行的 Imagenet 自定义数据集，包含归一化等操作。注： 嗯，我一直想做这件事很久了。我可能会同时再做一个 CIFAR 的自定义数据集。
• 一个适用于编码器-解码器网络（如 U-Net）的自定义数据集示例。注： 我的学生们已经为他们的项目编写了一个 U-Net 数据集，感兴趣的话可以查看 这里。

PyTorch 自定义数据集快速上手指南

本指南基于 pytorch-custom-dataset-examples 项目，旨在帮助开发者快速掌握如何在 PyTorch 中构建灵活的自定义 Dataset 类。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux, macOS 或 Windows
• Python: 3.6 及以上版本
• 核心依赖:
  • torch: PyTorch 深度学习框架
  • torchvision: 包含常用图像变换和数据集工具
  • pandas: 用于处理 CSV 数据文件（可选，视需求而定）
  • Pillow (PIL): 用于图像读取和处理

安装命令：

推荐使用国内镜像源（如清华源）以加速安装：

pip install torch torchvision pandas pillow -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

构建自定义数据集的核心是继承 torch.utils.data.Dataset 类，并实现三个关键方法：__init__、__getitem__ 和 __len__。

1. 最简骨架代码

这是构建任何自定义数据集的基础模板：

from torch.utils.data.dataset import Dataset

class MyCustomDataset(Dataset):
    def __init__(self, ...):
        # 初始化逻辑：读取文件列表、设置变换等
        pass
        
    def __getitem__(self, index):
        # 根据索引加载单个样本（如图片）和标签
        # 必须返回 Tensor、numpy array、dict 或 list 等类型
        img = ... 
        label = ...
        return (img, label)

    def __len__(self):
        # 返回数据集样本总数
        return count 

2. 集成 Torchvision 变换

通常需要在数据加载时应用预处理（如裁剪、归一化）。推荐在 __init__ 中接收变换对象，并在 __getitem__ 中应用。

from torch.utils.data.dataset import Dataset
from torchvision import transforms

class MyCustomDataset(Dataset):
    def __init__(self, ..., transforms=None):
        self.transforms = transforms
        
    def __getitem__(self, index):
        # 假设从文件加载了原始图像数据
        data = load_image_from_file(index) 
        
        # 应用变换
        if self.transforms is not None:
            data = self.transforms(data)
            
        label = get_label(index)
        return (data, label)

    def __len__(self):
        return total_count
        
if __name__ == '__main__':
    # 定义变换组合
    transformations = transforms.Compose([
        transforms.CenterCrop(100), 
        transforms.ToTensor()
    ])
    
    # 实例化数据集
    custom_dataset = MyCustomDataset(..., transformations)

3. 结合 Pandas 读取 CSV

当数据路径和标签存储在 CSV 文件中时，可使用 pandas 进行解析。

import pandas as pd
import numpy as np
from PIL import Image
from torch.utils.data.dataset import Dataset
from torchvision import transforms

class CustomDatasetFromImages(Dataset):
    def __init__(self, csv_path, transforms=None):
        self.transforms = transforms
        # 读取 CSV (假设无表头：第 1 列为路径，第 2 列为标签)
        self.data_info = pd.read_csv(csv_path, header=None)
        self.image_arr = np.asarray(self.data_info.iloc[:, 0])
        self.label_arr = np.asarray(self.data_info.iloc[:, 1])
        self.data_len = len(self.data_info.index)

    def __getitem__(self, index):
        # 获取图片路径并打开
        single_image_name = self.image_arr[index]
        img_as_img = Image.open(single_image_name)

        # 应用变换
        if self.transforms is not None:
            img_as_tensor = self.transforms(img_as_img)
        else:
            # 若无外部变换，至少需转为 Tensor
            img_as_tensor = transforms.ToTensor()(img_as_img)

        # 获取标签
        single_image_label = self.label_arr[index]

        return (img_as_tensor, single_image_label)

    def __len__(self):
        return self.data_len

4. 配合 DataLoader 使用

定义好 Dataset 后，即可通过 DataLoader 进行批量加载和多进程读取：

from torch.utils.data import DataLoader

# 实例化数据集
dataset = MyCustomDataset(..., transforms=transformations)

# 创建 DataLoader
loader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=4)

# 遍历数据
for images, labels in loader:
    # 训练逻辑
    pass

注意：__getitem__ 返回的单个样本必须是可被自动堆叠的类型（如 Tensor、numpy 数组、数字、字典或列表），否则 DataLoader 会报错。