【pytorch（07）】数据加载器，构建数据类Dataset类、TensorDataset类，数据集加载案例

旭日东升的xu.

已于 2025-08-07 19:37:59 修改

阅读量873

点赞数 25

CC 4.0 BY-SA版权

分类专栏：深度学习文章标签： pytorch 人工智能 python

于 2025-08-07 19:34:08 首次发布

本文链接：https://blog.csdn.net/2302_76685892/article/details/150020800

深度学习专栏收录该内容

11 篇文章

订阅专栏

【pytorch（01）】CUDA、cuDNN、pytorch安装，神奇魔法解决网络问题

【pytorch（02）】Tensor(张量)概述、如何创建、常见属性，切换设备

【pytorch（03）】Tensor(张量)数据转换、常见操作(元素值运算、阿达玛积、相乘、形状操作)

【pytorch（04）】自动微分：基础概念，梯度计算，梯度上下文控制

【pytorch（05）】感知神经网络，构造人工神经元，深入神经网络（了解结构，如何构建），全连接神经网络

【pytorch（06）】全连接神经网络：基本组件认知，线性层、激活函数、损失函数、优化器

数据加载器

分数据集和加载器2个步骤~

1.构建数据类

Dataset类

Dataset是一个抽象类，是所有自定义数据集应该继承的基类。它定义了数据集必须实现的方法。

必须实现的方法

__len__: 返回数据集的大小
__getitem__: 支持整数索引，返回对应的样本

在 PyTorch 中，构建自定义数据加载类通常需要继承 torch.utils.data.Dataset 并实现以下几个方法：

_init_ 方法
用于初始化数据集对象：通常在这里加载数据，或者定义如何从存储中获取数据的路径和方法。

def __init__(self, data, labels):
    self.data = data
    self.labels = labels

_len_ 方法
返回样本数量：需要实现，以便 Dataloader加载器能够知道数据集的大小。

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

_getitem_ 方法
根据索引返回样本：将从数据集中提取一个样本，并可能对样本进行预处理或变换。

def __getitem__(self, index):
    sample = self.data[index]
    label = self.labels[index]
    return sample, label

如果你需要进行更多的预处理或数据变换，可以在 _getitem_ 方法中添加额外的逻辑。

整体参考代码如下：

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader


# 定义数据加载类
class CustomDataset(Dataset):
    def __init__(self, data, labels):
        """
        初始化数据集
        :data: 样本数据（例如，一个 NumPy 数组或 PyTorch 张量）
        :labels: 样本标签
        """
        self.data = data
        self.labels = labels

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

    def __getitem__(self, index):
        index = min(max(index, 0), len(self.data) - 1)
        sample = self.data[index]
        label = self.labels[index]
        return sample, label


def test001():
    # 简单的数据集准备
    data_x = torch.randn(666, 20, requires_grad=True, dtype=torch.float32)
    data_y = torch.randn(data_x.shape[0], 1, dtype=torch.float32)
    dataset = CustomDataset(data_x, data_y)
    # 随便打印个数据看一下
    print(dataset[0])


if __name__ == "__main__":
    test001()

TensorDataset类

TensorDataset是Dataset的一个简单实现，它封装了张量数据，适用于数据已经是张量形式的情况。

特点

简单快捷：当数据已经是张量形式时，无需自定义Dataset类
多张量支持：可以接受多个张量作为输入，按顺序返回
索引一致：所有张量的第一个维度必须相同，表示样本数量

源码：

class TensorDataset(Dataset):
    def __init__(self, *tensors):
        # size(0)在python中同shape[0]，获取的是样本数量
        # 用第一个张量中的样本数量和其他张量对比，如果全部相同则通过断言，否则抛异常
        assert all(tensors[0].size(0) == tensor.size(0) for tensor in tensors)
        self.tensors = tensors

    def __getitem__(self, index):
        return tuple(tensor[index] for tensor in self.tensors)

    def __len__(self):
        return self.tensors[0].size(0)

示例：

def test03():

    torch.manual_seed(0)

    # 创建特征张量和标签张量
    features = torch.randn(100, 5)  # 100个样本，每个样本5个特征
    labels = torch.randint(0, 2, (100,))  # 100个二进制标签

    # 创建TensorDataset
    dataset = TensorDataset(features, labels)

    # 使用方式与自定义Dataset相同
    print(len(dataset))  # 输出: 100
    print(dataset[0])  # 输出: (tensor([...]), tensor(0))

2.数据加载器

在训练或者验证的时候，需要用到数据加载器批量的加载样本。

DataLoader 是一个迭代器，用于从 Dataset 中批量加载数据。它的主要功能包括：

批量加载：将多个样本组合成一个批次。
打乱数据：在每个 epoch 中随机打乱数据顺序。
多线程加载：使用多线程加速数据加载。

创建DataLoader：

# 创建 DataLoader
dataloader = DataLoader(
    dataset,          # 数据集
    batch_size=10,    # 批量大小
    shuffle=True,     # 是否打乱数据
    num_workers=2     # 使用 2 个子进程加载数据
)

遍历：

# 遍历 DataLoader
# enumerate返回一个枚举对象（iterator），生成由索引和值组成的元组
for batch_idx, (samples, labels) in enumerate(dataloader):
    print(f"Batch {batch_idx}:")
    print("Samples:", samples)
    print("Labels:", labels)

案例：

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader


# 定义数据加载类
class CustomDataset(Dataset):
	#略......


def test01():
    # 简单的数据集准备
    data_x = torch.randn(666, 20, requires_grad=True, dtype=torch.float32)
    data_y = torch.randn(data_x.size(0), 1, dtype=torch.float32)
    dataset = CustomDataset(data_x, data_y)

    # 构建数据加载器
    data_loader = DataLoader(dataset, batch_size=8, shuffle=True)
    for i, (batch_x, batch_y) in enumerate(data_loader):
        print(batch_x, batch_y)
        break

if __name__ == "__main__":
    test01()

数据集加载案例

通过一些数据集的加载案例，真正了解数据类及数据加载器。

1.加载csv数据集

代码参考如下:

import torch
from torch.utils.data import Dataset, DataLoader
import pandas as pd


class MyCsvDataset(Dataset):
    def __init__(self, filename):
        df = pd.read_csv(filename)
        # 删除文字列
        df = df.drop(["学号", "姓名"], axis=1)
        # 转换为tensor
        data = torch.tensor(df.values)
        # 最后一列以前的为data，最后一列为label
        self.data = data[:, :-1]
        self.label = data[:, -1]
        self.len = len(self.data)

    def __len__(self):
        return self.len

    def __getitem__(self, index):
        idx = min(max(index, 0), self.len - 1)
        return self.data[idx], self.label[idx]


def test001():
    excel_path = r"./大数据答辩成绩表.csv"
    dataset = MyCsvDataset(excel_path)
    dataloader = DataLoader(dataset, batch_size=4, shuffle=True)
    for i, (data, label) in enumerate(dataloader):
        print(i, data, label)


if __name__ == "__main__":
    test001()

练习：上述示例数据构建器改成TensorDataset

def build_dataset(filepath):
    df = pd.read_csv(filepath)
    df.drop(columns=['学号', '姓名'], inplace=True)
    data = df.iloc[..., :-1]
    labels = df.iloc[..., -1]

    x = torch.tensor(data.values, dtype=torch.float)
    y = torch.tensor(labels.values)

    dataset = TensorDataset(x, y)

    return dataset


def test001():
    filepath = r"./大数据答辩成绩表.csv"
    dataset = build_dataset(filepath)
    dataloader = DataLoader(dataset, batch_size=4, shuffle=True)
    for i, (data, label) in enumerate(dataloader):
        print(i, data, label)

2.加载图片数据集

使用ImageFolder加载图片集

ImageFolder 会根据文件夹的结构来加载图像数据。它假设每个子文件夹对应一个类别，文件夹名称即为类别名称。例如，一个典型的文件夹结构如下：

root/
    class1/
        img1.jpg
        img2.jpg
        ...
    class2/
        img1.jpg
        img2.jpg
        ...
    ...

在这个结构中：

root 是根目录。
class1、class2 等是类别名称。
每个类别文件夹中的图像文件会被加载为一个样本。

ImageFolder构造函数如下：

torchvision.datasets.ImageFolder(root, transform=None, target_transform=None, is_valid_file=None)

参数解释

root：字符串，指定图像数据集的根目录。
transform：可选参数，用于对图像进行预处理。通常是一个 torchvision.transforms 的组合。
target_transform：可选参数，用于对目标（标签）进行转换。
is_valid_file：可选参数，用于过滤无效文件。如果提供，只有返回 True 的文件才会被加载。

import torch
from torchvision import datasets, transforms
import os
from torch.utils.data import DataLoader
from matplotlib import pyplot as plt

torch.manual_seed(42)

def load():
    path = os.path.join(os.path.dirname(__file__), 'dataset')
    print(path)

    transform = transforms.Compose([
        transforms.Resize((112, 112)),
        transforms.ToTensor()
    ])

    dataset = datasets.ImageFolder(path, transform=transform)
    dataloader = DataLoader(dataset, batch_size=1, shuffle=True)

    for x,y in dataloader:
        x = x.squeeze(0).permute(1, 2, 0).numpy()
        plt.imshow(x)
        plt.show()
        print(y[0])
        break


if __name__ == '__main__':
    load()

3.加载官方数据集

在 PyTorch 中官方提供了一些经典的数据集，如 CIFAR-10、MNIST、ImageNet 等，可以直接使用这些数据集进行训练和测试。

数据集：https://pytorch.org/vision/stable/datasets.html

常见数据集：

MNIST: 手写数字数据集，包含 60,000 张训练图像和 10,000 张测试图像。
CIFAR10: 包含 10 个类别的 60,000 张 32x32 彩色图像，每个类别 6,000 张图像。
CIFAR100: 包含 100 个类别的 60,000 张 32x32 彩色图像，每个类别 600 张图像。
COCO: 通用对象识别数据集，包含超过 330,000 张图像，涵盖 80 个对象类别。

torchvision.transforms 和 torchvision.datasets 是 PyTorch 中处理计算机视觉任务的两个核心模块，它们为图像数据的预处理和标准数据集的加载提供了强大支持。

transforms 模块提供了一系列用于图像预处理的工具，可以将多个变换组合成处理流水线。

datasets 模块提供了多种常用计算机视觉数据集的接口，可以方便地下载和加载。

参考如下：

import torch
from torch.utils.data import Dataset, DataLoader
import torchvision
from torchvision import transforms, datasets


def test():
    transform = transforms.Compose(
        [
            transforms.ToTensor(),
        ]
    )
    # 训练数据集
    data_train = datasets.MNIST(
        root="./data",
        train=True,
        download=True,
        transform=transform,
    )
    trainloader = DataLoader(data_train, batch_size=8, shuffle=True)
    for x, y in trainloader:
        print(x.shape)
        print(y)
        break

    # 测试数据集
    data_test = datasets.MNIST(
        root="./data",
        train=False,
        download=True,
        transform=transform,
    )
    testloader = DataLoader(data_test, batch_size=8, shuffle=True)
    for x, y in testloader:
        print(x.shape)
        print(y)
        break


def test006():
    transform = transforms.Compose(
        [
            transforms.ToTensor(),
        ]
    )
    # 训练数据集
    data_train = datasets.CIFAR10(
        root="./data",
        train=True,
        download=True,
        transform=transform,
    )
    trainloader = DataLoader(data_train, batch_size=4, shuffle=True, num_workers=2)
    for x, y in trainloader:
        print(x.shape)
        print(y)
        break
    # 测试数据集
    data_test = datasets.CIFAR10(
        root="./data",
        train=False,
        download=True,
        transform=transform,
    )
    testloader = DataLoader(data_test, batch_size=4, shuffle=False, num_workers=2)
    for x, y in testloader:
        print(x.shape)
        print(y)
        break


if __name__ == "__main__":
    test()
    test006()