解决pytorch中Dataloader读取数据太慢的问题-Toy模板网

这篇具有很好参考价值的文章主要介绍了解决pytorch中Dataloader读取数据太慢的问题。希望对大家有所帮助。如果存在错误或未考虑完全的地方，请大家不吝赐教，您也可以点击"举报违法"按钮提交疑问。

提示：文章写完后，目录可以自动生成，如何生成可参考右边的帮助文档

前言

最近在使用pytorch框架进行模型训练时遇到一个性能问题，即数据读取的速度远远大于GPU训练的速度，导致整个训练流程中有大部分时间都在等待数据发送到GPU，在资源管理器中呈现出CUDA使用率周期性波动，且大部分时间都是在等待数据加载。

一、造成的原因

其实从前言中就可以知道，造成这样的原因可以理解为：GPU的算力与数据加载速度之间的不匹配。

二、查找不匹配的原因

本人使用的GPU为RTX3060，数据集为cifar10，使用的模型为VGG，显然这张显卡对于这个任务来说是绰绰有余的，无论是显存还是算力。
但是几经测试发现，数据从内存送到GPU的速度实在是太慢了，去百度了很久都没有很好的解决。那回到这个问题的本身，既然是数据加载导致的性能差，那问题就出在pytorch的dataset和dataloader中。

在dataset中，会将数据从磁盘读入内存中，如果启用了dataloader中的pin_memory，就会让数据常驻内存，同时设置num_workers还能实现多进程读取数据，但即使设置了这些，数据加载速度依然没有质的提升。

博主发现，dataset中的transform是导致性能慢的一个原因，dataset中有个函数为__getitem__，每获取一个数据就会让这个数据过一次transform。相信大家都写过如下的代码：

transform = transforms.Compose([
	transforms.RandomHorizontalFlip(),
    transforms.ToTensor(),
    transforms.Normalize([0.4914, 0.4822, 0.4465], [0.5, 0.5, 0.5])
])

但是这里的ToTensor和Normalize完全没必要没读一次数据都处理一次，可以在数据加载到内存的时候就直接全部处理完，这样每个数据只需要经历一次ToTensor和Normalize，这会大大提高数据读取速度，大家可以自行测试一次ToTensor和Normalize所需要的时间，还是非常大的。

在训练的过程中，相信大家也写过如下代码：

for x, y in dataloader:
	x, y = x.cuda(), y.cuda()

经过博主测试，将数据发送到GPU也是非常耗时的，那为什么不一次性全部加载到GPU里面呢？当然前提是你的GPU显存够大。

三、解决方法

以上分析可以总结为两点：

数据的预处理有一部分可以提前对全部数据做一遍；
如果显存足够，可以将数据全部加载到GPU中。

基于此，我们可以重载类，这里以pytorch自带的cifar10为例：

class CUDACIFAR10(CIFAR10):
    def __init__(
            self, 
            root: str, 
            train: bool = True,
            to_cuda: bool = True,
            half: bool = False,
            pre_transform: Optional[Callable] = None,
            transform: Optional[Callable] = None, 
            target_transform: Optional[Callable] = None, 
            download: bool = False) -> None:

        super().__init__(root, train, transform, target_transform, download)
        if pre_transform is not None:
            self.data = self.data.astype("float32")
            for index in range(len(self)):
                """
                ToTensor的操作会检查数据类型是否为uint8, 如果是, 则除以255进行归一化, 这里data提前转为float,
                所以手动除以255.
                """
                self.data[index] = pre_transform(self.data[index]/255.0).numpy().transpose((1, 2, 0))
                self.targets[index] = torch.Tensor([self.targets[index]]).squeeze_().long()
                if to_cuda:
                    self.targets[index] = self.targets[index].cuda()
            self.data = torch.Tensor(self.data).permute((0, 3, 1, 2))
            if half:
                self.data = self.data.half()
            if to_cuda:
                self.data = self.data.cuda()

    def __getitem__(self, index: int) -> Tuple[Any, Any]:
        """
        Args:
            index (int): Index

        Returns:
            tuple: (image, target) where target is index of the target class.
        """
        img, target = self.data[index], self.targets[index]

        if self.transform is not None:
            img = self.transform(img)
        if self.target_transform is not None:
            target = self.target_transform(target)
        return img, target

to_cuda为True就会让数据全部加载到GPU中，后续就不需要写x, y = x.cuda(), y.cuda()了。
pre_transform就是让所有数据提前进行的处理，例如使用ToTensor和Normalize，后续调用时不会再进行这些处理。
transform为后续调用时会进行的处理，一般就是一些随机处理过程。

在博主的测试过程中发现，解决了以上问题后，一个epoch只要2秒就能完成，而平时需要15秒，并且任务管理器中的CUDA几乎全程拉满。唯一的代价就是显存占用更高了，这何尝不是一种空间换时间。

四、使用方法

这里直接粘贴我为这个类写的注释

- 使用pytorch自带的CIFAR10时, 每读取一个数据都会调用一次transforms, 其中ToTensor()和Normalize()会消耗巨大的时间
  如果你的数据集非常的大, 那么一个epoch将会花费非常多的时间用于读取数据, 如果还要将数据送入GPU, 那么时间将会继续增加。

- 一般的写法如下：

    for epoch in range(epochs):
        for x, y in dataloader:
            x, y = x.cuda(), y.cuda()

  如果你的数据集很大并且GPU算力很强, 那么读取数据并发送的GPU将会成为性能瓶颈。

- CUDACIFAR10是专门针对pytorch的CIFAR10进行优化的, 使用的前提是你的显存足够的大, 至少8GB, 且读取数据已经是性能瓶颈。
  CUDACIFAR10的参数与CIFAR10非常相似, 新增的参数为：
    to_cuda: bool, 是否将数据集常驻GPU, default=True
    half: bool, 进一步降低数据所占据的显存, 在混合精度训练时使用, 否则可能存在意外(例如梯度值overflow)
    pre_transform: 传入一个transforms, 如果不为None, 则会在初始化数据时直接对所有数据进行对应的转换, 在后续调用时将
                    不会使用pre_transform. 可以将ToTensor()和Normalize()作为pre_transform, 这会大幅度减少读取时间.

- CUDACIFAR10的用法如下:

    pre_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize([0.4914, 0.4822, 0.4465], [0.5, 0.5, 0.5])
    ])
    dataset = CUDACIFAR10(..., to_cuda=True, pre_transform=pre_transform, ...)
    dataloader = Dataloader(dataset, ..., pin_memory=False, num_workers=0, ...)
    ...
    for epoch in range(epochs):
        for x, y in dataloader:
            # 不需要写x, y = x.cuda(), y.cuda(), 除非to_cuda=False
            ...

- 使用CUDACIFAR10需要注意如果启用了to_cuda, 那么Dataloader不能启用pin_memory, pin_memory是将数据常驻内存, 这会产生冲突.
  同时num_workers=0.

- 如果参数to_cuda=False, pre_transform=None, 那么该类与CIFAR10用法完全一致.