PyTorch|Dataset与DataLoader使用、构建自定义数据集


一、Dataset与DataLoader

PyTorch提供的两个常用数据API:

  • torch.utils.data.Dataset:用于处理单个训练样本,读取数据特征、size、标签等,并且包括数据转换等;
  • torch.utils.data.DataLoader:DataLoader在Dataset周围重载一个可迭代对象,以便轻松访问样本。

官方案例: Fashion-MNIST数据集
torchvision:torch的一个视觉库,将torchvision中的datasets导入进来,就能获得其中的各种数据集

FashionMNIST图像存储在目录img_dir中,标签存储在CSV文件annotations_file中

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()
)

对上述数据集进行可视化:

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()

二、自定义Dataset类

  • 构建自定义的Dataset类,需要继承TensorFlow的官方dataset类
  • 自定义Dataset类必须实现三个函数:__init__,__len__和__getitem__

pytorch中的dataset类是在pytorch的torch下的utils之下的data文件夹里有一个dataset.py
在这里插入图片描述

(一)__init__函数

包含图像、注释文件和两个转换:

  • annotations_file:标注文件
  • img_dir:图像目录
def __init__(self, annotations_file, img_dir, transform=None, target_transform=None):
    self.img_labels = pd.read_csv(annotations_file) #标签存储在CSV文件annotations_file中
    self.img_dir = img_dir #FashionMNIST图像存储在目录img_dir中
    self.transform = transform #图像转换
    self.target_transform = target_transform

(二)__len__函数

返回数据集的样本数(就是img_labels的长度)

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

(三)__getitem_\函数

输入索引index,getitem函数从数据集中加载并返回对应index的一个样本:

def __getitem__(self, idx):
		#img_labels的第index行第0列标注了对应的照片文件名称
    img_path = os.path.join(self.img_dir, self.img_labels.iloc[idx, 0])
    image = read_image(img_path) #使用read_image将图像转换为张量
    label = self.img_labels.iloc[idx, 1] #从self中的csv数据中检索相应的标签
    #调用转换函数
    if self.transform: 
        image = self.transform(image)
    if self.target_transform:
        label = self.target_transform(label)
    return image, label #返回张量图像和相应的标签

(四)全部代码

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

三、将单个样本组成minibatch(DataLoader)

(一)PyTorch的DataLoader源码

1、DataLoader的参数

DataLoader通常是在torch.utils.data下
在这里插入图片描述
常用的参数有:

  • dataset(数据集):需要提取数据的数据集,Dataset对象
  • batch_size(批大小):每一次装载样本的个数,int型
  • shuffle:是否打乱数据顺序
  • sampler(Sampler, optional): 自定义从数据集中取样本的策略,如果指定这个参数,那么shuffle必须为False
  • num_workers:进行数据加载时使用单个进程还是多进程进行加载,多进程意为加载速度更快,一般默认为0,表示使用主进程进行加载
  • collate_fn (callable, optional): 将一个list的sample组成一个mini-batch的函数,一般用于对于一个batch进行后处理
  • pin_memory (bool, optional): 如果设置为True,那么data loader将会在返回它们之前将tensors拷贝到CUDA中的固定内存(CUDA pinned memory)中
  • drop_last:当样本数不能被batchsize整除时, 是否舍弃最后一批数据
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=True)

2、init函数

主要做了三件事:构建sampler、构建batch_sampler、构建collate_fn

定义属性:
在这里插入图片描述
如果设置了自定义的sampler然后又设置了shuffle=true,这种情况是没有意义的:
(shuffle是官方自定义的一个随机sampler)
在这里插入图片描述
设置了batch_sampler的情况下,就不需要设置batch_size、shuffle、sampler和drop_last了:
在这里插入图片描述
如果没有设置sampler,则先判断数据集类型,如果使用的是map-style(else逻辑),就根据是否设置shuffle来选择pytorch内置的sampler:
在这里插入图片描述
设置了batch_size但是没有设置batch_sampler时,会使用内置的BatchSampler:
在这里插入图片描述
如果没有设置collate_fn,就判断auto_collation是否设置(auto_collation是根据batch_sampler是否是None来设置的,如果batch_sampler不是none,auto_collation就是true),default_collate是将batch作为输入,batch输出,并没有对数据做额外处理:
在这里插入图片描述

3、iter函数

iter函数返回的是get_iterator的值:
在这里插入图片描述
get_iterator根据num_workers的设置选择对应的内置DataLoaderIter:
在这里插入图片描述

所以可知,iter函数最终返回的是一个dataloaderiter对象,以SingleProcessDataLoaderIter为例,类里有next_data函数:
在这里插入图片描述
SingleProcessDataLoaderIter类是继承了BaseDataLoaderIter类,BaseDataLoaderIter类中的next函数就是使用了子类中的next_data:
在这里插入图片描述

(二)使用DataLoader遍历

根据上述源码分析,就可以对dataloader去迭代iter之后调用next函数来获得每一批次的数据:

  • 通过DataLoader实现对于数据集的遍历,每次遍历会得到一个batch的数据,这里设置batch_size为64:
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=True)
  • iter函数将train_dataloader变成一个迭代器,使用next函数可以以此从迭代器中生成一个一个的批次:
# 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}")

在这里插入图片描述由于batch_size=64,因此最终返回的Feature batch shape以及Labels batch shape均为64。


参考:
PyTorch官方文档:Datasets & DataLoaders
5、深入剖析PyTorch DataLoader源码

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值