Pytorch Dataloader加速

在多卡训练中,CPU生成数据的速度可能无法满足GPU处理数据的需求,导致GPU利用率低。本文介绍了两种方法来加速Pytorch Dataloader:1) 调整线程数,但效果有限且可能在Windows上引发问题;2) 使用Data Prefetcher,通过创建单独的Stream实现数据加载和计算的并行,提高运行速度,但会增加显存占用。记得在使用Data Prefetcher时设置Dataloader的pin_memory为True。

摘要生成于 C知道 ,由 DeepSeek-R1 满血版支持, 前往体验 >

Pytorch Dataloader加速

在进行多卡训练的时候,经常会出现GPU利用率上不来的情况,无法发挥硬件的最大实力。 造成这种现象最有可能的原因是,CPU生成数据的能力,已经跟不上GPU处理数据的能力。


方法一


常见的方法为修改Dataloader里面的线程数量,利用多线程技术提高数据生产能力,但是这种方法提速并不是特别明显。

train_loader = DataLoader(dataset, batch_size,shuffle=True, num_worker=4)

而且windows机器上,num_worker大于0时,有时会出现卡死的情况,这应该是pytorch的bug,因此不是特别建议这种方法。
不过这种方法最简单,还是可以尝试一下更改线程数能否缓解你遇到的问题。nun_worker一般设置为处理器的物理线程数,不宜过大,因为会导致额外的线程开销。

方法二


本文主要介绍第二种方法,也就是Data Prefetcher,最早见于NVIDIA APEX

这里我把代码抠出来了,删除掉了一些不必要的注释,可以将其复用到自己的项目里来。

import torch

class data_prefetcher():
    def __init__(self, loader):
        self.loader = iter(loader)
        self.stream = torch.cuda.Stream()
        self.mean = torch.tensor([0.485 * 255, 0.456 * 255, 0.406 * 255]).cuda().view(1,3,1,1)
        self.std = torch.tensor([0.229 * 255, 0.224 * 255, 0.225 * 255]).cuda().view(1,3,1,1)
        self.preload()

    def preload(self):
        try:
            self.next_input, self.next_target = next(self.loader)
        except StopIteration:
            self.next_input = None
            self.next_target = None
            return
    
        with torch.cuda.stream(self.stream):
            self.next_input = self.next_input.cuda(non_blocking=True)
            self.next_target = self.next_target.cuda(non_blocking=True)
            self.next_input = self.next_input.float()
            self.next_input = self.next_input.sub_(self.mean).div_(self.std)

    def next(self):
        torch.cuda.current_stream().wait_stream(self.stream)
        input = self.next_input
        target = self.next_target
        if input is not None:
            input.record_stream(torch.cuda.current_stream())
        if target is not None:
            target.record_stream(torch.cuda.current_stream())
        self.preload()
        return input, target

首先我们来看初始化函数,在初始化函数中,会直接调用preload,所以当这个对象初始化时,就会生成第一份的输入数据。

核心逻辑也就在预加载函数preload中,其中第13行是从原来的dataloader中取数,这一步和常规数据加载没有差别。有差别的是第19行,这里出现了Stream的概念。

一般来说,CUDA程序默认都运行在同一个Stream上,因此CPU->GPU,GPU->GPU以及GPU->CPU的一系列计算都是在同一个Stream里面串行运行的。 深度学习一般流程是先从dataloader中取数,这里是内存->CPU的运算,然后执行to_device操作,让数据从CPU->GPU,再是GPU->GPU的神经网络计算。

代码19行,使得data_prefetecher这个类是单独运行在一个Stream上的,因此它让数据加载和神经网络计算可以并行执行,也就加速了整体的运行速度。这样做带来的负面结果就是GPU同时在做两项任务,所以显存占用会增加。

这里不知道解释清楚没有,建议去看一下原作者的回答link

另外,重要的是,使用这个方法的时候一定要将Dataloader里面的pin_memory设置为True。

使用方法如下,非常简单,改造前是从dataloader里取数,改造后是将dataloader包在prefetecher里面,从prefetecher里面取数。

train_loader = DataLoader(dataset, batch_size,shuffle=True, num_worker=4,pin_memory=True)
prefetcher = data_prefetcher(train_loader)
input, target = prefetcher.next()

while input is not None:
    ## 
    前后向计算...
    ###
    input, target = prefetcher.next()
### 回答1: PyTorch DataLoader是一个用于批量加载数据的工具,它可以帮助用户在训练模型时高效地加载和处理大规模数据集。DataLoader可以根据用户定义的批量大小、采样方法、并行加载等参数来自动将数据集分成小批量,并且可以在GPU上并行加载数据以提高训练效率。 使用DataLoader需要先定义一个数据集对象,然后将其传递给DataLoader。常用的数据集对象包括PyTorch自带的Dataset类和用户自定义的数据集类。在DataLoader中可以指定批量大小、是否打乱数据、并行加载等参数。 下面是一个示例代码: ```python import torch from torch.utils.data import Dataset, DataLoader class MyDataset(Dataset): def __init__(self): self.data = torch.randn(100, 10) self.label = torch.randint(0, 2, size=(100,)) def __getitem__(self, index): return self.data[index], self.label[index] def __len__(self): return len(self.data) dataset = MyDataset() dataloader = DataLoader(dataset, batch_size=10, shuffle=True, num_workers=2) for data, label in dataloader: print(data.shape, label.shape) ``` 在上面的示例中,我们定义了一个自己的数据集类MyDataset,并将其传递给DataLoader。然后指定了批量大小为10,打乱数据,使用2个进程来并行加载数据。在循环中,每次从DataLoader中取出一个批量的数据和标签,并输出它们的形状。 ### 回答2: PyTorchDataLoader是一个用于加载数据的实用工具。它可以帮助我们高效地加载和预处理数据,以供深度学习模型使用。 DataLoader有几个重要参数。首先是dataset,它定义了我们要加载的原始数据集。PyTorch提供了几种内置的数据集类型,也可以自定义数据集。数据集可以是图片、文本、音频等。 另一个重要参数是batch_size,它定义了每个批次中加载的数据样本数量。这是非常重要的,因为深度学习模型通常需要在一个批次上进行并行计算。较大的批次可以提高模型的训练速度,但可能需要更多的内存。 DataLoader还支持多线程数据加载。我们可以使用num_workers参数来指定并行加载数据的线程数。这可以加快数据加载的速度,特别是当数据集很大时。 此外,DataLoader还支持数据的随机打乱。我们可以将shuffle参数设置为True,在每个轮次开始时随机重新排序数据。这对于训练深度学习模型非常重要,因为通过在不同轮次中提供不同样本的顺序,可以增加模型的泛化能力。 在使用DataLoader加载数据后,我们可以通过迭代器的方式逐批次地获取数据样本。每个样本都是一个数据批次,包含了输入数据和对应的标签。 总的来说,PyTorchDataLoader提供了一个简单而强大的工具,用于加载和预处理数据以供深度学习模型使用。它的灵活性和可定制性使得我们可以根据实际需求对数据进行处理,并且能够高效地并行加载数据,提高了训练的速度。 ### 回答3: PyTorchDataLoader是一个用于数据加载和预处理的实用程序类。它可以帮助我们更有效地加载和处理数据集,并将其用于训练和评估深度学习模型。 DataLoader的主要功能包括以下几个方面: 1. 数据加载:DataLoader可以从不同的数据源中加载数据,例如文件系统、内存、数据库等。它接受一个数据集对象作为输入,该数据集对象包含实际的数据和对应的标签。DataLoader可以根据需要将数据集分成小批量加载到内存中,以减少内存占用和加速训练过程。 2. 数据预处理:DataLoader可以在加载数据之前对数据进行各种预处理操作,包括数据增强、标准化、裁剪和缩放等。这些预处理操作可以提高模型的泛化能力和训练效果。 3. 数据迭代:DataLoader将数据集划分为若干个小批量,并提供一个可迭代的对象,使得我们可以使用for循环逐个访问这些小批量。这种迭代方式使得我们能够更方便地按批次处理数据,而无需手动编写批处理循环。 4. 数据并行加载:DataLoader支持在多个CPU核心上并行加载数据,以提高数据加载的效率。它使用多线程和预读取的机制,在一个线程中预先加载数据,而另一个线程处理模型的训练或推理过程。 总之,PyTorchDataLoader是一个方便且高效的工具,帮助我们更好地管理和处理数据集。它可以加速深度学习模型的训练过程,并提供了一种简单而灵活的数据加载和迭代方式。使用DataLoader可以让我们更专注于模型的设计和调优,而无需过多关注数据的处理和加载细节。
评论 2
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

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

抵扣说明:

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

余额充值