PyTorch中 Datasets & DataLoader 的介绍-阿里云开发者社区

PyTorch中 Datasets & DataLoader 的介绍

2024-05-14 137

版权

本文内容由阿里云实名注册用户自发贡献，版权归原作者所有，阿里云开发者社区不拥有其著作权，亦不承担相应法律责任。具体规则请查看《阿里云开发者社区用户服务协议》和《阿里云开发者社区知识产权保护指引》。如果您发现本社区中有涉嫌抄袭的内容，填写侵权投诉表单进行举报，一经查实，本社区将立刻删除涉嫌侵权内容。

简介： PyTorch中 Datasets & DataLoader 的介绍

前言

用于处理数据样本的代码可能很快就会变得混乱且难以维护。理想情况下，为了获得更好的可读性和模块化，我们希望处理数据集的代码与模型训练代码分离。

PyTorch 提供了两个非常有用的数据集处理类：

torch.utils.data.Dataset：存储样本及其相应的标签，PyTorch还提供了不少自带的数据集。
torch.utils.data.DataLoader：围绕Dataset包装一个可迭代对象，以便轻松访问样本。

PyTorch 提供了许多预加载的数据集（例如：FashionMNIST），它们是 torch.utils.data.Dataset的子类并实现特定于特定数据的函数。我们可以用它们来对模型进行原型设计和基准测试。这些数据集可以分为：图像数据集、文本数据集和音频数据集。

1、加载数据集

现在我们来展示一下如何从 TorchVision 加载 Fashion-MNIST 数据集。Fashion-MNIST由60000个训练样本和10000个测试样本组成。每个样本包含一个 28x28 灰度图像和来自10个类别之一的关联标签。

我们使用以下参数加载 FashionMNIST数据集：

root 是存储训练/测试数据的路径
train 指定训练或测试数据集
download = True 如果root目录下没有数据，则从网上下载数据

transform 和 target_transform 指定特征和标签转换

import torch
from torch.utils.data import Dataset
from torchvision import datasets
from torchvision.transforms import ToTensor
import matplotlib.pyplot as plt
 
 
training_data = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor()
)
 
test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor()
)

2、遍历并可视化数据集

我们可以用索引来访问数据集中的样本，用 matplotlib 可视化图形样本。

labels_map = {
    0: "T-Shirt",
    1: "Trouser",
    2: "Pullover",
    3: "Dress",
    4: "Coat",
    5: "Sandal",
    6: "Shirt",
    7: "Sneaker",
    8: "Bag",
    9: "Ankle Boot",
}
figure = plt.figure(figsize=(8, 8))
cols, rows = 3, 3
for i in range(1, cols * rows + 1):
    sample_idx = torch.randint(len(training_data), size=(1,)).item()
    img, label = training_data[sample_idx]
    figure.add_subplot(rows, cols, i)
    plt.title(labels_map[label])
    plt.axis("off")
    plt.imshow(img.squeeze(), cmap="gray")
plt.show()

3、从本地文件创建自己的数据集

自定义 Dataset 类必须实现三个函数：

__init__：在实例化 Dataset 对象时运行一次。我们初始化包含图像的目录、注释文件和 transform 与 target_transform。
__len__：以 len(dataset)的方式获取 dataset 中包含的样本数
__getitem__：加载并返回给定索引 idx 处的数据集样本。基于索引，它识别图像在磁盘上的位置，使用read_image将其转换为Tensor，从self.img_labels中的CSV数据中检索相应的标签，调用它们的转换函数（如果适用），并以元组的形式返回Tensor图像和相应的标签。

import os
import pandas as pd
from torchvision.io import read_image
 
class CustomImageDataset(Dataset):
    def __init__(self, annotations_file, img_dir, transform=None, target_transform=None):
        self.img_labels = pd.read_csv(annotations_file)
        self.img_dir = img_dir
        self.transform = transform
        self.target_transform = target_transform
 
    def __len__(self):
        return len(self.img_labels)
 
    def __getitem__(self, idx):
        img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0])
        image = read_image(img_path)
        label = self.img_labels.iloc[idx, 1]
        if self.transform:
            image = self.transform(image)
        if self.target_transform:
            label = self.target_transform(label)
        return image, label

4、使用 DataLoader 准备数据以进行训练

Dataset 检索数据集的特征并一次标记一个样本。在训练模型时，我们通常希望以小批量（mini batch）方式传递样本，在每个epoch重新整理数据以减少模型过拟合，并使用Python的多线程来加速数据检索。

DataLoader 是一个可迭代的对象。它通过一个简单的API为我们抽象了这种复杂性需求。

from torch.utils.data import DataLoader
 
train_dataloader = DataLoader(training_data, batch_size=64, shuffle=True)
test_dataloader = DataLoader(test_data, batch_size=64, shuffle=False)

5、遍历 DatasetLoader

我们已将该数据集加载到 DataLoader中，并且可以根据需要迭代数据集。

下面的每次迭代都会返回一批 train_features 和 train_labels（分别包含 batch_size=64个特征和标签）。

# Display image and label.
train_features, train_labels = next(iter(train_dataloader))
print(f"Feature batch shape: {train_features.size()}")
print(f"Labels batch shape: {train_labels.size()}")
img = train_features[0].squeeze()
label = train_labels[0]
plt.imshow(img, cmap="gray")
plt.show()
print(f"Label: {label}")

参考： PyTorch研习社

PyTorch中 Datasets & DataLoader 的介绍

前言

1、加载数据集

2、遍历并可视化数据集

3、从本地文件创建自己的数据集

4、使用 DataLoader 准备数据以进行训练

5、遍历 DatasetLoader

热门文章

最新文章

相关课程

相关电子书

推荐镜像

热门

活动广场

任务中心

开发者评测

高校计划

乘风者计划

训练营

阿里云MVP

话题

直播

下载

镜像站

技术资料

插件

PyTorch中 Datasets & DataLoader 的介绍

前言

1、加载数据集

2、遍历并可视化数据集

3、从本地文件创建自己的数据集

4、使用 DataLoader 准备数据以进行训练

5、遍历 DatasetLoader

热门文章

最新文章

相关课程

相关电子书

推荐镜像