在深度学习的应用中,数据是模型训练和评估的基础。有效地管理和加载数据对于构建高效的神经网络模型至关重要。PyTorch,作为一个广泛使用的深度学习框架,提供了一套强大的工具来帮助开发者处理数据。Dataset
和DataLoader
是PyTorch中用于构建数据管道的核心组件,它们使得数据的加载、预处理和批量化变得更加简单和高效。
Dataset
是一个抽象类,它定义了数据集的接口,允许开发者以统一的方式访问数据。通过实现__len__
方法,Dataset
提供了数据集的大小信息,而__getitem__
方法则允许按索引获取数据集中的单个样本。这种设计使得Dataset
可以被看作是一个动态数组,可以根据需要获取数据集中的任何部分。
DataLoader
则是一个迭代器,它负责将Dataset
中的数据以批次(batch)的形式输出。通过DataLoader
,我们可以轻松地控制批次的大小,设置随机采样或顺序采样策略,并在必要时使用多进程来加速数据的加载。此外,DataLoader
还提供了一个collate_fn
参数,允许开发者自定义数据整理函数,以处理复杂的数据类型或特殊的数据需求。
使用Dataset
和DataLoader
,我们可以构建灵活且高效的数据管道,为模型的训练和评估提供强有力的支持。这不仅简化了数据管理的工作,还提高了代码的可读性和可维护性。在接下来的内容中,我们将详细介绍如何使用这两个工具来构建自己的数据管道,以及如何通过自定义函数来满足特定的数据处理需求。
自学笔记,欢迎交流;
0. DataSet和DataLoader介绍
Pytorch 使用 Dataset和DataLoader这两个工具来构建数据管道
DataSet定义了数据集的内容,类似于列表的数据结构,具有确定的长度,能够索引获取长度;
DataLoader定义了按batch加载数据集的方法,是一个可迭代对象,每次迭代输出一个batch的数据;
DataLoader能够控制batch的大小,batch中元素的采样方法,以及将batch结果整理成模型所需要的输入形式的方法,并能够使用多进程读取数据;
通常只需要实现dataset的__len__和__getitem__方法,就可以构建自己数据集,并且用默认数据管道进行加载;
复杂数据,需要设计dataloader中的collate_fn方法将一个批次的数据整理成模型需要的输入形式;
DataSet和DataLoader原理
1.1 如何获取一个batch数据 : 总样 —— 抽样 —— 取样值 —— 聚合
(假设数据集的特征和标签分别为张量X和Y,数据集可以表示为(X, Y), batch为M)
- 先要获得数据的长度N。假设N=1000
- 然后从0—N-1的范围中抽样出M个数,假设M=4,那么拿到的结果就是 indices = [1,4,8,9](4个随机数)
- 接着我们去数据集中取这M个数分别对应的下标,这里拿到一个元组列表 samples = [(X[1], Y[1]), (X[4], Y[4]), (X[8], Y[8]), (X[9], Y[9])]
- 最后将结果整理成两个张量作为输出,故最后的结果是两个张量 batch = (features, labels)
其中 features = torch.stack([[X][1], X[4], X[8], X[9]]) ; labels = torch.stack([Y[1], Y[4], Y[8], Y[9]]
1.2 DataSet和DataLoader的功能分工
步骤1获取长度是DataSet中的__len__方法获取的
步骤2从0—N-1的范围中抽样M个数是DataLoader的sampler和batch_sampler参数指定的;
sampler 参数指定抽样方法,一般不需要设置,程序默认在DataLoader的参数shuffle=True时采用随机抽样,shuffle=False时采用顺序抽样;
batch_sampler参数将多个抽样的元素整理成一个列表,一般无需用户设置,默认方法在DataLoader的参数drop_last=True时会丢弃数据集最后一个长度不能被batch大小整除的批次,在drop_last=False时保留最后一个批次。
步骤3的核心逻辑就是根据下标取值,由DataSet中的__getitem__实现;
步骤4是由DataLoader的参数collate_fn指定,一般情况下无需设置;
一般使用方法如下
import torch
from torch.utils.data import TensorDataset,Dataset,DataLoader
from torch.utils.data import RandomSampler,BatchSampler
ds = TensorDataset(torch.randn(1000,3),
torch.randint(low=0,high=2,size=(1000,)).float()) # 步骤1
dl = DataLoader(ds,batch_size=4,drop_last = False) # 步骤2 3 4
features,labels = next(iter(dl)) # 取数据
print("features = ",features ) # 打印查看
print("labels = ",labels )
DataLoader内部调用方式步骤拆解
# step1: 确定数据集长度 (Dataset的 __len__ 方法实现)
ds = TensorDataset(torch.randn(1000,3),
torch.randint(low=0,high=2,size=(1000,)).float())
print("n = ", len(ds)) # len(ds)等价于 ds.__len__()
# step2: 确定抽样indices (DataLoader中的 Sampler和BatchSampler实现)
sampler = RandomSampler(data_source = ds)
batch_sampler = BatchSampler(sampler = sampler,
batch_size = 4, drop_last = False)
for idxs in batch_sampler:
indices = idxs
break
print("indices = ",indices)
# step3: 取出一批样本batch (Dataset的 __getitem__ 方法实现)
batch = [ds[i] for i in indices] # ds[i] 等价于 ds.__getitem__(i)
print("batch = ", batch)
# step4: 整理成features和labels (DataLoader 的 collate_fn 方法实现)
def collate_fn(batch): # 聚合函数
features = torch.stack([sample[0] for sample in batch])
labels = torch.stack([sample[1] for sample in batch])
return features,labels
features,labels = collate_fn(batch)
print("features = ",features)
print("labels = ",labels)
1.3 Dataset和DataLoader的核心源码
import torch
class Dataset(object):
def __init__(self):
pass
def __len__(self): # 获取长度信息
raise NotImplementedError
def __getitem__(self,index): # 获取索引信息
raise NotImplementedError
class DataLoader(object):
def __init__(self,dataset,batch_size,collate_fn = None,shuffle = True,drop_last = False):
self.dataset = dataset
self.sampler =torch.utils.data.RandomSampler if shuffle else \
torch.utils.data.SequentialSampler # 抽样方式 随机/顺序
self.batch_sampler = torch.utils.data.BatchSampler
self.sample_iter = self.batch_sampler(
self.sampler(self.dataset), # 索引取值
batch_size = batch_size,drop_last = drop_last) # 要不要最后一组batch
self.collate_fn = collate_fn if collate_fn is not None else \ # 聚合函数
torch.utils.data._utils.collate.default_collate
def __next__(self):
indices = next(iter(self.sample_iter))
batch = self.collate_fn([self.dataset[i] for i in indices])
return batch
def __iter__(self):
return self
测试
class ToyDataset(Dataset):
def __init__(self,X,Y):
self.X = X
self.Y = Y
def __len__(self):
return len(self.X)
def __getitem__(self,index):
return self.X[index],self.Y[index]
X,Y = torch.randn(1000,3),torch.randint(low=0,high=2,size=(1000,)).float()
ds = ToyDataset(X,Y)
dl = DataLoader(ds,batch_size=4,drop_last = False)
features,labels = next(iter(dl))
print("features = ",features )
print("labels = ",labels )
2. 使用DataSet创建数据集
Dataset创建数据集常用的方法有:
- 使用 torch.utils.data.TensorDataset 根据Tensor创建数据集(numpy的array,Pandas的DataFrame需要先转换成Tensor)。
- 使用 torchvision.datasets.ImageFolder 根据图片目录创建图片数据集。
- 继承 torch.utils.data.Dataset 创建自定义数据集。
此外,还可以通过
- torch.utils.data.random_split 将一个数据集分割成多份,常用于分割训练集,验证集和测试集。
- 调用Dataset的加法运算符(+)将多个数据集合并成一个数据集。
2.1 根据Tensor创建数据集
import numpy as np
import torch
from torch.utils.data import TensorDataset,Dataset,DataLoader,random_split
# 根据Tensor创建数据集
from sklearn import datasets
iris = datasets.load_iris() # 导入鸢尾花数据
ds_iris = TensorDataset(torch.tensor(iris.data),torch.tensor(iris.target)) # 转成tensor数据
# 分割成训练集和预测集
n_train = int(len(ds_iris)*0.8)
n_val = len(ds_iris) - n_train
ds_train,ds_val = random_split(ds_iris,[n_train,n_val])
print(type(ds_iris))
print(type(ds_train))
# 使用DataLoader加载数据集
# 最简单的给两个参数,一个数据集,一个batch_size大小
dl_train,dl_val = DataLoader(ds_train,batch_size = 8),DataLoader(ds_val,batch_size = 8)
for features,labels in dl_train:
print(features,labels)
break
# 演示加法运算符(`+`)的合并作用
ds_data = ds_train + ds_val
print('len(ds_train) = ',len(ds_train))
print('len(ds_valid) = ',len(ds_val))
print('len(ds_train+ds_valid) = ',len(ds_data))
print(type(ds_data))
2.2 根据图片目录创建图片数据集
import numpy as np
import torch
from torch.utils.data import DataLoader
from torchvision import transforms,datasets
#演示一些常用的图片增强操作
from PIL import Image
img = Image.open('./data/cat.jpeg')
img
# 随机数值翻转
transforms.RandomVerticalFlip()(img)
#随机旋转
transforms.RandomRotation(45)(img)
# 定义图片增强操作
transform_train = transforms.Compose([
transforms.RandomHorizontalFlip(), #随机水平翻转
transforms.RandomVerticalFlip(), #随机垂直翻转
transforms.RandomRotation(45), #随机在45度角度内旋转
transforms.ToTensor() #转换成张量
]
)
transform_valid = transforms.Compose([ # 评估数据不做增强
transforms.ToTensor()
]
)
# 根据图片目录创建数据集
def transform_label(x):
return torch.tensor([x]).float()
ds_train = datasets.ImageFolder("./eat_pytorch_datasets/cifar2/train/",
transform = transform_train,target_transform= transform_label)
ds_val = datasets.ImageFolder("./eat_pytorch_datasets/cifar2/test/",
transform = transform_valid,
target_transform= transform_label)
print(ds_train.class_to_idx)
# 使用DataLoader加载数据集
dl_train = DataLoader(ds_train,batch_size = 50,shuffle = True)
dl_val = DataLoader(ds_val,batch_size = 50,shuffle = True)
for features,labels in dl_train:
print(features.shape)
print(labels.shape)
break
2.3 创建自定义数据集
下面我们通过另外一种方式,即继承 torch.utils.data.Dataset 创建自定义数据集的方式来对 cifar2构建 数据管道。
from pathlib import Path
from PIL import Image
# 自定义数据集,继承Dataset
class Cifar2Dataset(Dataset):
# 传入文件目录与数据增强方式
def __init__(self,imgs_dir,img_transform):
# 这里是图像的目录,目录中是图片及其编号
self.files = list(Path(imgs_dir).rglob("*.jpg"))
self.transform = img_transform
def __len__(self,):
# 获取长度
return len(self.files)
def __getitem__(self,i):
# 逐个操作,根据文件的编号获取相对应的数据与标签信息
file_i = str(self.files[i])
img = Image.open(file_i)
tensor = self.transform(img)
label = torch.tensor([1.0]) if "1_automobile" in file_i else torch.tensor([0.0])
return tensor,label
train_dir = "./eat_pytorch_datasets/cifar2/train/"
test_dir = "./eat_pytorch_datasets/cifar2/test/"
# 定义图片增强
transform_train = transforms.Compose([
transforms.RandomHorizontalFlip(), #随机水平翻转
transforms.RandomVerticalFlip(), #随机垂直翻转
transforms.RandomRotation(45), #随机在45度角度内旋转
transforms.ToTensor() #转换成张量
]
)
transform_val = transforms.Compose([
transforms.ToTensor()
]
)
ds_train = Cifar2Dataset(train_dir,transform_train)
ds_val = Cifar2Dataset(test_dir,transform_val)
dl_train = DataLoader(ds_train,batch_size = 50,shuffle = True)
dl_val = DataLoader(ds_val,batch_size = 50,shuffle = True)
for features,labels in dl_train:
print(features.shape)
print(labels.shape)
break
3. 使用DataLoader加载数据集
DataLoader能够控制batch的大小,batch中元素的采样方法,以及将batch结果整理成模型所需输入形式的方法,并且能够使用多进程读取数据。
DataLoader的函数签名如下。
DataLoader(
dataset,
batch_size=1,
shuffle=False,
sampler=None,
batch_sampler=None,
num_workers=0,
collate_fn=None,
pin_memory=False,
drop_last=False,
timeout=0,
worker_init_fn=None,
multiprocessing_context=None,
)
一般情况下,我们仅仅会配置 dataset, batch_size, shuffle, num_workers,pin_memory, drop_last这六个参数,
有时候对于一些复杂结构的数据集,还需要自定义collate_fn函数,其他参数一般使用默认值即可。
DataLoader除了可以加载我们前面讲的 torch.utils.data.Dataset 外,还能够加载另外一种数据集 torch.utils.data.IterableDataset。
和Dataset数据集相当于一种列表结构不同,IterableDataset相当于一种迭代器结构。它更加复杂,一般较少使用。
- dataset : 数据集
- batch_size: 批次大小
- shuffle: 是否乱序
- sampler: 样本采样函数,一般无需设置。
- batch_sampler: 批次采样函数,一般无需设置。
- num_workers: 使用多进程读取数据,设置的进程数。
- collate_fn: 整理一个批次数据的函数。
- pin_memory: 是否设置为锁业内存。默认为False,锁业内存不会使用虚拟内存(硬盘),从锁业内存拷贝到GPU上速度会更快。
- drop_last: 是否丢弃最后一个样本数量不足batch_size批次数据。
- timeout: 加载一个数据批次的最长等待时间,一般无需设置。
- worker_init_fn: 每个worker中dataset的初始化函数,常用于 IterableDataset。一般不使用。
参考文章 实践教程|源码级理解Pytorch中的Dataset和DataLoader - CV技术指南(公众号) - 博客园