Pytorch加载数据的Dateset类和DataLoader类

Pytorch提供了Dataset类和DataLoader类专门用于处理数据,他们既可以加载Pytorch预置的数据集,也可以加载自定义数据集。其中数据集类Dataset负责存储样本以及他们对应的标签;数据加载类DataLoader负责迭代访问数据集中的样本。

Dataset

  • 映射型数据集

    继承自Dataset类,表示一个从索引到样本的映射(索引可以不是整数),这样我们就可以方便地通过dataset[idx] 来访问指定得到索引的样本。这也是目前最常见的数据集类型。映射型数据集必须实现 getitem ()函数,其负责根据指定的key返回对应的样本。一般还会实现 len() 用于返回数据集的大小。

  • 迭代型数据集

    继承自 Dataset,表示可迭代的数据集,它可以通过 iter(dataset)以数据流(steam)的形式访问,适用于访问大数据集或者远程服务器上产生的数据,迭代数据集必须实现__iter__()函数,用于返回一个样本迭代器(iteror).
    多线程访问数据集

python 复制代码
from torch.utils.data import IterableDataset, DataLoader

class MyIterableDataset(IterableDataset):

    def __init__(self, start, end):
        super(MyIterableDataset).__init__()
        assert end > start
        self.start = start
        self.end = end

    def __iter__(self):
        return iter(range(self.start, self.end))

ds = MyIterableDataset(start=3, end=7) # [3, 4, 5, 6]
# Single-process loading
print(list(DataLoader(ds, num_workers=0)))
# Directly doing multi-process loading
print(list(DataLoader(ds, num_workers=2)))

# output
[tensor([3]), tensor([4]), tensor([5]), tensor([6])]

[tensor([3]), tensor([3]), tensor([4]), tensor([4]), tensor([5]), tensor([5]), tensor([6]), tensor([6])]

创建一个自定义的映射型数据集

python 复制代码
import os
import pandas as pd
from torchvision.io import read_image
from torch.utils.data import Dataset

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

可以看到,我们实现了 init ()、len () 和 getitem() 三个函数,其中:

  • init() 初始化数据集参数,这里设置了图像的存储目录、标签(通过读取标签 csv 文件)以及样本和标签的数据转换函数;
  • len() 返回数据集中样本的个数;
  • getitem() 映射型数据集的核心,根据给定的索引 idx 返回样本。这里会根据索引从目录下通过 read_image 读取图片和从 csv 文件中读取图片标签,并且返回处理后的图像和标签。

DataLoader

前面的数据集Dataset类提供了一种按照索引访问样本的方式,不过在实际训练模型时,我们都需要先将数据集切分为很多的mini-batches,然后按照批次将样本送入模型,并且循环这一过程,每一个完整遍历所有样本成为一个epoch。

python 复制代码
from torch.utils.data import DataLoader
from torchvision import datasets
from torchvision.transforms import ToTensor

training_data = datasets.FashionMNIST(
    root="data",
    train=True,
    download=True,
    transform=ToTensor()
)

test_data = datasets.FashionMNIST(
    root="data",
    train=False,
    download=True,
    transform=ToTensor()
)

train_dataloader = DataLoader(training_data, batch_size=64, shuffle=True)
test_dataloader = DataLoader(test_data, batch_size=64, shuffle=True)

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]
print(img.shape)
print(f"Label: {label}")
相关推荐
子燕若水4 小时前
Unreal Engine 5中的AI知识
人工智能
极限实验室5 小时前
Coco AI 实战(一):Coco Server Linux 平台部署
人工智能
杨过过儿5 小时前
【学习笔记】4.1 什么是 LLM
人工智能
巴伦是只猫5 小时前
【机器学习笔记Ⅰ】13 正则化代价函数
人工智能·笔记·机器学习
大千AI助手5 小时前
DTW模版匹配:弹性对齐的时间序列相似度度量算法
人工智能·算法·机器学习·数据挖掘·模版匹配·dtw模版匹配
AI生存日记5 小时前
百度文心大模型 4.5 系列全面开源 英特尔同步支持端侧部署
人工智能·百度·开源·open ai大模型
LCG元6 小时前
自动驾驶感知模块的多模态数据融合:时序同步与空间对齐的框架解析
人工智能·机器学习·自动驾驶
why技术6 小时前
Stack Overflow,轰然倒下!
前端·人工智能·后端
烛阴6 小时前
简单入门Python装饰器
前端·python