Pytorch 加速读取数据之 prefetch_factor

最新推荐文章于 2025-09-05 08:30:44 发布
原创 最新推荐文章于 2025-09-05 08:30:44 发布 · 1.6w 阅读 · 14 ·
CC 4.0 BY-SA版权
版权声明：本文为博主原创文章，遵循 CC 4.0 BY-SA 版权协议，转载请附上原文出处链接和本声明。

博客分享了提高GPU训练效率的两种方法。一是利用NVIDIA DALI模块加速数据预处理，通过并行处理减少训练延迟。二是升级PyTorch到1.7以上版本，利用新增的`prefetch_factor`参数，预先加载更多批次的数据，显著缩短数据读取时间，提高显卡利用率。作者实测将加载时间从3s降至0.5s，效果明显。
部署运行你感兴趣的模型镜像

最近使用四卡训练图片，发现总有卡的效率突变到0，大致就是在读取以及处理数据了
怎么能加速读取呢？一个方法是用NVIDIA的DALI模块，可以加速，具体可以参考 英伟达DALI加速技巧：让数据预处理速度比原生PyTorch快4倍
主要就是通过并行训练和预处理过程，减少了延迟及训练时间
在这里插入图片描述

但是今天我发现一个更简单的方法
就是升级pytorch到1.7以上，目前是1.8.1，最好的1.8.1吧
为什么会这么说呢，因为在dataloader中加入了一个参数 prefetch_factor，这个就是提前加载多少个batch的数据，具体更改看github ，具体说如下，现在默认prefetch_factor =2 ，就是意味着预先加载 prefetch 2 * num_workers 个data

fix #40604
Add parameter to Dataloader to configure the per-worker prefetch number.
Before this edit, the prefetch process always prefetch 2 * num_workers data items, this commit help us make this configurable, e.x. you can specify to prefetch 10 * num_workers data items

我在我这里试了一下，升级以后，有从3s到0.5s，速度还是很明显的，而且显卡利用率都上来了。大家可以自己试试~

您可能感兴趣的与本文相关的镜像

PyTorch 2.5

PyTorch 2.5

PyTorch
Cuda

PyTorch 是一个开源的 Python 机器学习库，基于 Torch 库，底层由 C++ 实现，应用于人工智能领域，如计算机视觉和自然语言处理

如何大幅加速 PyTorch 模型训练
人工智能领域专家
07-18 819
3. 优化 DataLoader 速度训练深度学习模型有时就像看着油漆慢慢变干。迭代过程可能缓慢到让你忍不住用头撞键盘，眼睁睁看着终端上每个 epoch 像蜗牛般爬行，这时你难免会想：有没有更好的方法？别担心！我又带来了一篇文章，分享一些让你的深度学习工作流程更加敏捷的技巧/窍门。我将分享一些直接而有用的代码调整，帮助你更有效地利用计算资源。我们还将深入探讨一系列经过行业验证的技巧、窍门和调整，这些方法可以从你的硬件中榨取每一分性能，让你少花时间等待，多花时间进行更快的迭代。
pytorch 数据的读取
qq_40107571的博客
09-01 1955
pytorch 数据装载器dataloader ,以及超分辨率dataset的常规定义
提速80%：PyTorch多线程数据处理实战指南
最新发布
gitblog_01094的博客
09-05 362
你是否还在忍受模型训练时CPU数据加载的漫长等待？是否遇到过GPU利用率忽高忽低的"锯齿状"曲线？本文将通过PyTorch的`DataLoader`多线程机制，手把手教你解决数据加载瓶颈，让GPU真正吃饱跑满。读完本文你将掌握： - 多线程数据加载的核心参数调优 - 硬件资源与`num_workers`的匹配公式 - 内存锁定与预取策略的实战配置 - 常见性能陷阱与解决方案 ## 数据加载的并行...
pytorch中dataloader的prefetch_factor出错
horsetaill的博客
02-18 3746
可以看到报错信息是：prefetch_factor这个属性只有在num_workers大于0的时候才能被定义。但是我去看了看源码，很明显train源码中定义没有问题啊！然后我在github上的pytorch的issue中发现了有人提出了疑问。这是跑示例遇到的最后一个坑！
DataLoader是每一个worker加载一个batch数据，并不是所有的worker一起加载一个batch数据
qq_41818256的博客
11-26 1892
DataLoader是每一个worker加载一个batch数据，并不是所有的worker一起加载一个batch数据。
DataLoader PyTorch 主要参数的含义
子燕若水的博客
08-22 4251
DataLoader类是一个用于从数据集（dataset）中加载数据，并以迭代器（iterator）的形式返回数据样本（data samples）的工具¹²。pin_memory将数据样本保存在锁页内存，这里的内存指的是cpu的内存还是gpu的内存？
Pytorch DataLoader prefetch_factor pin_memory
码匀的博客
10-25 5766
dataset (Dataset) ：所加载的数据集。 batch_size (int, optional) ：每个批次加载多少个样本。 shuffle (bool, optional) ：如果为 True，则每个epoch 都会再打乱数据。 sampler (Sampler or Iterable, optional) ：定义了如何从样本采样的策略。可以是任何实现了 __len__的迭代器。 batch_sampler (Sampler or Iterable, optional) ：与sampler类.
pytorch 数据预加载
撒旦先生的博客
12-04 2877
本文介绍一个工具，它包装，接收该类的一个实例loader，启动一个线程t，创建一个队列qt将loader中的数据预加载到队列q中, 以在模型计算时也能启动启动数据加载程序, 节省数据加载时间。"""** 包装 torch.utils.data.DataLoader, 接收该类的一个实例 loader, 启动一个线程 t, 创建一个队列 qt 将 loader 中的数据预加载到队列 q 中, 以在模型计算时也能启动启动数据加载程序, 节省数据加载时间。
PyTorch DataLoader 参数详解
qq_43069203的博客
04-08 1408
使用DataLoaderbatch_size：控制每个批次的样本数；shuffle：决定数据顺序是否随机打乱；：通过多进程加速数据加载；pin_memory：用于加快 GPU 数据传输；drop_last：确保批次样本数一致；collate_fn：自定义 batch 拼接逻辑。在训练 Stable Diffusion（SD）模型时，通过自定义Dataset类整合图片和 JSON 数据，并合理设置DataLoader的参数，可以显著提升数据加载效率和训练稳定性。
pytorch学习笔记1 - 数据集读取
fa1c4的blog
11-16 1679
前言 要开始搞data science了, 最近在考虑一些AI结合security的事情, 还比较迷茫, 不过应该总能找到出路的, 所以先不管学这个有什么意义了, 直接开干技术. pytorch同TensorFlow一样, 是data science标配, 所以有必要好好学习一下 torch包结构 查看torch包的成员 dir(torch) ['AVG', 'AggregationType', 'AnyType', 'Argument', 'ArgumentSpec', 'BFloat16Sto
02_pytorch中的DataLoader
chumingqian的博客
08-31 2770
sampler 是用来定义取batch方法的一个函数或者类，返回的是一个迭代器。我们可以看下自带的RandomSampler类中最重要的iter函数。
pytorch conv2d计算shape,prefetch_factor提速，提高显卡利用率
power的专栏
05-31 641
W×WF×F。
Pytorch Dataloader参数补充篇（collate_fn、pin_memory、dorp_last、prefetch_factor、persistent_workers）
gy77
08-07 2657
如果不为None，那么在设定随机种子后，数据加载前，这个函数会在每个子进程上调用，并把子进程id作为输入（一个在[0, num_workers - 1]的整数），默认为None。数据预处理过程中出现了问题:如果数据预处理函数(如collate_fn或自定义的数据转换函数)出现了错误或死循环,可能会导致数据加载器无法获取下一个数据批次。多进程/多线程数据加载时出现了死锁或竞争条件:当使用多进程或多线程加载数据时,如果出现了死锁或竞争条件,也可能导致数据加载器无法获取下一个数据批次。
pytorch使用data_prefetcher提升数据读取速度
Mr_health的博客
02-09 5795
直接给出代码： class data_prefetcher(): def __init__(self, loader): #loader 1：real #loader 2：fake self.stream = torch.cuda.Stream() self.loader = iter(loader) self.preload() def preload(self): try:
pytorch cv2 pil读取图片测试模型，数据加变量Variable,在GPU上加速处理
m0_37192554的博客
11-26 2811
def cv2img_process(img): ####分类图片预处理， img = one_img(img)#add pad white resize 224 assert (img.shape[0]==224 and img.shape[1]==224),"img not resize 224" mean = [0.485, 0.456, 0.406]...
Pytorch之DataLoader参数说明
qq_44066464的博客
04-12 866
功能：组合数据集和采样器(规定提取样本的方法)，并提供对给定数据集的可迭代对象。通俗一点，就是把输进来的数据集，按照一个想要的规则（采样器）把数据划分好，同时让它是一个可迭代对象（可以循环提取数据，方便后面程序使用）。
【pytorch】给训练踩踩油门-- Pytorch 加速数据读取
Shwan_ma的博客
12-01 4732
本文转载于 https://zhuanlan.zhihu.com/p/80695364 训练模型的时候有时候会发现显卡的占用一直跑不满，会很浪费，往往是因为IO瓶颈导致的训练速度降低。 本文可以从以下几个方面进行对模型加速： 一, prefetch_generator 使用 prefetch_generator 库在后台加载下一 batch 的数据。 安装： pip install prefet...
pytorch dataloader_简单两步加速PyTorch里的Dataloader
weixin_39518002的博客
11-25 1358
PyTorch中通过Dataloader加载图片，使用十分方便。但当加载图片较多并且需要做较多变换时，加载的速度很慢，会出现加载数据过慢（即使已经使用了多个worker），GPU空闲等待数据加载的情况。这篇文章就和大家分享一下怎么给Dataloader提提速。1、读取jpg图片加速加载数据慢的一个很重要的原因的读取图片本身就比较慢，尤其是jpg图片解码更是耗时。如果你是读取jpg图片的话建议使用j...
pytorch数据加载-DataLoader
君莫笑
04-28 1893
torch.data
pytorch如何用多线程的方式加载数据集
08-28
<think>我们正在讨论PyTorch中如何使用多线程加载数据集。根据引用内容，主要涉及`DataLoader`中的`num_workers`参数。 引用[1]和[3]都提到了`DataLoader`的`num_workers`参数用于设置多线程加载数据。引用[4]的标题提到了Dataset多线程加载，但具体内容似乎与DataLoader相关。 关键点： 1. 在PyTorch中，多线程加载数据是通过`DataLoader`的`num_workers`参数实现的。 2. `num_workers`表示使用多少个子进程来加载数据。设置为0表示只在主进程中加载（没有多线程），大于0则表示使用多个子进程。 步骤： 1. 定义数据集（继承自`torch.utils.data.Dataset`）。 2. 创建`DataLoader`实例，并设置`num_workers`参数。 注意： - 在Windows系统中，多进程加载可能会遇到问题，因为Windows使用spawn而不是fork来创建进程。因此，在Windows下，需要将主要代码放在`if __name__ == '__main__':`中，以避免多进程问题。 - 设置`num_workers`的数量通常根据CPU核心数来定，一般设置为CPU核心数或2倍，但具体需要根据实际情况调整。 示例代码（基于引用[3]）： ```python from torch.utils.data import Dataset, DataLoader # 自定义数据集类，继承自Dataset class MyDataset(Dataset): def __init__(self, data, targets): self.data = data self.targets = targets def __len__(self): return len(self.data) def __getitem__(self, idx): sample = self.data[idx] label = self.targets[idx] return sample, label # 假设我们有一些数据和标签 data = ... # 你的数据 targets = ... # 你的标签 # 创建数据集实例 dataset = MyDataset(data, targets) # 创建DataLoader，设置num_workers为2（表示使用2个子进程加载数据） dataloader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=2) ``` 引用[1]中提供的torchvision.datasets加载MNIST的示例也使用了`num_workers=2`。 另外，引用[4]提到“Dataset多线程加载数据，而不是在DataLoader”，但根据其内容，似乎还是在讨论DataLoader的配置。实际上，在PyTorch中，多线程加载是在DataLoader中实现的，而不是在Dataset内部。Dataset类本身是单线程的，但是DataLoader通过多个worker（每个worker是一个进程）来并行调用Dataset的`__getitem__`方法，从而加速数据加载。 因此，总结如下： - 使用`DataLoader`并设置`num_workers`大于0即可实现多线程（实际上是多进程）加载数据。 - 每个worker进程会独立地加载数据，然后将批处理数据通过进程间通信传递给主进程。 注意事项： - 当使用多进程加载时，Dataset对象会被复制到每个worker进程中。因此，如果数据集很大，可能会占用较多内存。 - 在Unix-like系统（如Linux）上，由于使用fork创建进程，所以通常不会出现问题。但在Windows上，由于使用spawn，需要确保Dataset的定义是可以被pickle的，并且将主要代码放在main块中。 示例（适用于Windows）： ```python if __name__ == '__main__': # 创建DataLoader等代码 dataloader = DataLoader(dataset, batch_size=32, shuffle=True, num_workers=2) # 训练循环等 ``` 性能优化： - 如果数据加载是瓶颈（比如从磁盘读取大量小文件），可以增加`num_workers`的数量。 - 但注意，过多的worker可能会因为进程间通信的开销而降低性能，也可能导致内存不足。 因此，用户需要根据硬件配置和数据集特点调整`num_workers`的值。 引用文献： [^1]: 使用torchvision.datasets加载MNIST数据集，并设置num_workers=2 [^3]: DataLoader的参数num_workers用于多线程加载 [^4]: 标题提到Dataset多线程加载，但内容与DataLoader相关 相关问题： 1. 如何确定PyTorch中DataLoader的num_workers的最佳值？ 2. 在Windows系统下使用PyTorch多进程加载数据有哪些注意事项？ 3. PyTorch多进程加载数据时如何避免内存泄漏？ 4. 除了num_workers，还有哪些方法可以优化PyTorch数据加载速度？</think>在 PyTorch 中，使用多线程加载数据集主要通过 `DataLoader` 的 `num_workers` 参数实现。以下是详细方法和步骤： ### 1. 使用 `DataLoader` 的 `num_workers` 参数（推荐方法） 这是最常用的多线程加载方式，通过设置 `num_workers` 参数指定并行加载数据的子进程数量： ```python from torch.utils.data import DataLoader, Dataset # 1. 创建自定义数据集类（继承自Dataset） class MyDataset(Dataset): def __init__(self, data, targets): self.data = data self.targets = targets def __len__(self): return len(self.data) def __getitem__(self, idx): return self.data[idx], self.targets[idx] # 2. 实例化数据集 dataset = MyDataset(data, targets) # 3. 创建DataLoader并启用多线程加载 dataloader = DataLoader( dataset, batch_size=32, # 每批数据量 shuffle=True, # 是否打乱数据顺序 num_workers=4, # 关键参数：使用4个子进程加载数据 pin_memory=True # 加速GPU数据传输（如果使用GPU） ) ``` **关键参数说明**： - `num_workers`：指定并行加载数据的子进程数量（设为0表示主进程单线程加载） - 最佳值通常为CPU核心数（如4核CPU设为4）或2倍核心数 - `pin_memory=True` 可加速GPU数据传输（当使用CUDA时） ### 2. 使用 `torchvision.datasets` 内置数据集 对于常见数据集（如MNIST、CIFAR），可直接使用内置方法： ```python import torchvision.datasets as datasets import torchvision.transforms as transforms # 1. 定义数据预处理 transform = transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.5,), (0.5,)) ]) # 2. 加载数据集（自动下载） trainset = datasets.MNIST( root='./data', # 存储路径 train=True, # 训练集 download=True, # 自动下载 transform=transform # 数据预处理 ) # 3. 创建多线程DataLoader trainloader = DataLoader( trainset, batch_size=64, shuffle=True, num_workers=4, # 使用4个线程 pin_memory=True ) ``` ### 3. 高级优化技巧 1. **预取数据**（Prefetching）： ```python dataloader = DataLoader( dataset, num_workers=4, prefetch_factor=2 # 每个worker预取2个批次 ) ``` 2. **持久化worker**（避免重复创建进程）： ```python dataloader = DataLoader( dataset, num_workers=4, persistent_workers=True # 保持worker进程活跃 ) ``` 3. **自定义多线程加载逻辑**（高级用法）： ```python from torch.utils.data import DataLoader, Dataset import concurrent.futures class ParallelDataset(Dataset): def __init__(self, data): self.data = data self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=4) def __getitem__(self, idx): # 在线程池中执行加载任务 future = self.executor.submit(self._load_item, idx) return future.result() def _load_item(self, idx): # 实际数据加载逻辑 return self.data[idx] ``` ### 注意事项 1. **Windows系统限制**： - 需将主要代码放在 `if __name__ == '__main__':` 中 - 避免使用 `fork` 启动方式（PyTorch默认使用 `spawn`） 2. **资源管理**： - `num_workers` 过大可能导致内存溢出 - 监控CPU利用率（目标：80-90%） 3. **数据顺序保证**： - 当 `shuffle=True` 时，每个epoch会自动重置采样器 - 使用 `worker_init_fn` 可为每个worker设置独立随机种子 ```python def set_worker_seed(worker_id): worker_seed = torch.initial_seed() % 2**32 numpy.random.seed(worker_seed) random.seed(worker_seed) dataloader = DataLoader( dataset, num_workers=4, worker_init_fn=set_worker_seed ) ``` ### 性能对比 | 配置方式 | 数据加载速度 | CPU利用率 | 内存占用 | |---------|------------|----------|---------| | `num_workers=0` | 基准速度 | 20-30% | 低 | | `num_workers=4` | 3-5倍加速 | 70-90% | 中等 | | `num_workers=8` | 可能下降 | 100%+ | 高 | > **最佳实践**：从 `num_workers=CPU核心数` 开始测试，逐步增加直到性能不再提升[^1][^3]。
评论 1
添加红包

请填写红包祝福语或标题

个

红包个数最小为10个

元

红包金额最低5元

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

抵扣说明：

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

余额充值