聊聊Pytorch的dataloader

点击上方“机器学习与生成对抗网络”&＃xff0c;关注星标

获取有趣、好玩的前沿干货&＃xff01;

来源&＃xff1a;知乎—Mario 编辑 人工智能前沿讲习

地址&＃xff1a;https://zhuanlan.zhihu.com/p/117270644

为啥突然要写一下pytorch的dataloader呢&＃xff0c;首先来说说事情的来龙去脉。

起初&＃xff0c;我最开始单独训练一个网络来完成landmark点回归任务和分类任务&＃xff0c;训练的数据是txt格式&＃xff0c;在训练之前对数据进行分析&＃xff0c;发现分类任务中存在严重的数据样本不均衡的问题&＃xff0c;那么我事先针对性的进行数据采样均衡操作&＃xff0c;重新得到训练和测试的txt数据和标签&＃xff0c;保证了整个训练和测试数据的样本均衡性。由于我的整个项目是检测&＃43;点回归&＃43;分类&＃xff0c;起初检测和点回归&＃43;分类是分两步实现的&＃xff0c;检测是通过读取XML格式来进行训练&＃xff0c;现在要统一整个项目的训练和测试过程&＃xff0c;要将点回归&＃43;分类的训练测试过程也按照读取XML格式来进行&＃xff0c;那么就遇到一个问题&＃xff0c;如何针对性的去给样本偏少的样本进行均衡&＃xff0c;由于在dataset类中&＃xff0c;返回的图像和标签都是针对每个index返回一个结果&＃xff0c;在dataset类中进行操作似乎不太可行&＃xff0c;那么就想到在dataloader中进行操作&＃xff0c;通过dataloader中的参数sample来完成针对性采样。

还有一个问题是关于num_workers的设置&＃xff0c;因为我有对比过&＃xff0c;在我的单机RTX 2080Ti上和八卡服务器TITAN RTX上(仅使用单卡&＃xff0c;其它卡有在跑其它任务)&＃xff0c;使用相同的num_workers&＃xff0c;在单机上的训练速度反而更快&＃xff0c;于是猜想可能和CPU或者内存有关系&＃xff0c;下面会具体分析。

首先来看下下dataloader中的各个参数的含义。

类的定义为&＃xff1a;torch.utils.data.DataLoader &＃xff0c;其中包含的参数有&＃xff1a;

torch.utils.data.DataLoader(dataset, batch_size&＃61;1, shuffle&＃61;False, sampler&＃61;None, \batch_sampler&＃61;None, num_workers&＃61;0, collate_fn&＃61;None, pin_memory&＃61;False, \drop_last&＃61;False, timeout&＃61;0, worker_init_fn&＃61;None, multiprocessing_context&＃61;None)


dataset&＃xff1a;定义的dataset类返回的结果。

batchsize&＃xff1a;每个bacth要加载的样本数&＃xff0c;默认为1。

shuffle&＃xff1a;在每个epoch中对整个数据集data进行shuffle重排&＃xff0c;默认为False。

sample&＃xff1a;定义从数据集中加载数据所采用的策略&＃xff0c;如果指定的话&＃xff0c;shuffle必须为False&＃xff1b;batch_sample类似&＃xff0c;表示一次返回一个batch的index。

num_workers&＃xff1a;表示开启多少个线程数去加载你的数据&＃xff0c;默认为0&＃xff0c;代表只使用主进程。

collate_fn&＃xff1a;表示合并样本列表以形成小批量的Tensor对象。

pin_memory&＃xff1a;表示要将load进来的数据是否要拷贝到pin_memory区中&＃xff0c;其表示生成的Tensor数据是属于内存中的锁页内存区&＃xff0c;这样将Tensor数据转义到GPU中速度就会快一些&＃xff0c;默认为False。

drop_last&＃xff1a;当你的整个数据长度不能够整除你的batchsize&＃xff0c;选择是否要丢弃最后一个不完整的batch&＃xff0c;默认为False。

注&＃xff1a;这里简单科普下pin_memory&＃xff0c;通常情况下&＃xff0c;数据在内存中要么以锁页的方式存在&＃xff0c;要么保存在虚拟内存(磁盘)中&＃xff0c;设置为True后&＃xff0c;数据直接保存在锁页内存中&＃xff0c;后续直接传入cuda&＃xff1b;否则需要先从虚拟内存中传入锁页内存中&＃xff0c;再传入cuda&＃xff0c;这样就比较耗时了&＃xff0c;但是对于内存的大小要求比较高。

下面针对num_workers&＃xff0c;sample和collate_fn分别进行说明&＃xff1a;

01

设置num_workers

pytorch中dataloader一次性创建num_workers个子线程&＃xff0c;然后用batch_sampler将指定batch分配给指定worker&＃xff0c;worker将它负责的batch加载进RAM&＃xff0c;dataloader就可以直接从RAM中找本轮迭代要用的batch。如果num_worker设置得大&＃xff0c;好处是寻batch速度快&＃xff0c;因为下一轮迭代的batch很可能在上一轮/上上一轮...迭代时已经加载好了。坏处是内存开销大&＃xff0c;也加重了CPU负担&＃xff08;worker加载数据到RAM的进程是进行CPU复制&＃xff09;。如果num_worker设为0&＃xff0c;意味着每一轮迭代时&＃xff0c;dataloader不再有自主加载数据到RAM这一步骤&＃xff0c;只有当你需要的时候再加载相应的batch&＃xff0c;当然速度就更慢。num_workers的经验设置值是自己电脑/服务器的CPU核心数&＃xff0c;如果CPU很强、RAM也很充足&＃xff0c;就可以设置得更大些&＃xff0c;对于单机来说&＃xff0c;单跑一个任务的话&＃xff0c;直接设置为CPU的核心数最好。

02

定义sample&＃xff1a;&＃xff08;假设dataset类返回的是&＃xff1a;data, label&＃xff09;

from torch.utils.data.sampler import WeightedRandomSampler
## 如果label为1&＃xff0c;那么对应的该类别被取出来的概率是另外一个类别的2倍
weights &＃61; [2 if label &＃61;&＃61; 1 else 1 for data, label in dataset]
sampler &＃61; WeightedRandomSampler(weights,num_samples&＃61;10, replacement&＃61;True)
dataloader &＃61; DataLoader(dataset, batch_size&＃61;16, sampler&＃61;sampler)


PyTorch中提供的这个sampler模块&＃xff0c;用来对数据进行采样。默认采用SequentialSampler&＃xff0c;它会按顺序一个一个进行采样。常用的有随机采样器&＃xff1a;RandomSampler&＃xff0c;当dataloader的shuffle参数为True时&＃xff0c;系统会自动调用这个采样器&＃xff0c;实现打乱数据。这里使用另外一个很有用的采样方法&＃xff1a;WeightedRandomSampler&＃xff0c;它会根据每个样本的权重选取数据&＃xff0c;在样本比例不均衡的问题中&＃xff0c;可用它来进行重采样。replacement用于指定是否可以重复选取某一个样本&＃xff0c;默认为True&＃xff0c;即允许在一个epoch中重复采样某一个数据。

03

定义collate_fn

def detection_collate(batch):"""Custom collate fn for dealing with batches of images that have a differentnumber of associated object annotations (bounding boxes).Arguments:batch: (tuple) A tuple of tensor images and lists of annotationsReturn:A tuple containing:1) (tensor) batch of images stacked on their 0 dim2) (list of tensors) annotations for a given image are stacked on0 dim"""targets &＃61; []imgs &＃61; []for sample in batch:imgs.append(sample[0])targets.append(torch.FloatTensor(sample[1]))return torch.stack(imgs, 0), targets


使用dataloader时加入collate_fn参数&＃xff0c;即可合并样本列表以形成小批量的Tensor对象&＃xff0c;如果你的标签不止一个的话&＃xff0c;还可以支持自定义&＃xff0c;在上述方法中再额外添加对应的label即可。

data_loader &＃61; torch.utils.data.DataLoader(dataset, args.batch_size, num_workers&＃61;args.num_workers, sampler&＃61;sampler, shuffle&＃61;False,     collate_fn&＃61;detection_collate, pin_memory&＃61;True, drop_last&＃61;True)
参考链接&＃xff1a;
https://pytorch.org/docs/stable/data.html?highlight&＃61;dataloader#torch.utils.data.DataLoader
https://discuss.pytorch.org/t/guidelines-for-assigning-num-workers-to-dataloader
猜您喜欢&＃xff1a;
超100篇&＃xff01;CVPR 2020最全GAN论文梳理汇总&＃xff01;拆解组新的GAN&＃xff1a;解耦表征MixNMatchStarGAN第2版&＃xff1a;多域多样性图像生成
附下载 | 《可解释的机器学习》中文版附下载 |《TensorFlow 2.0 深度学习算法实战》附下载 |《计算机视觉中的数学方法》分享《基于深度学习的表面缺陷检测方法综述》《零样本图像分类综述: 十年进展》《基于深度神经网络的少样本学习综述》