PyTorch：加速数据加载

Posted 2023-02-16

【中文标题】PyTorch：加速数据加载

【英文标题】：PyTorch: Speed up data loading

【发布时间】：2020-08-07 03:43:54

【问题描述】：

我正在使用 densenet121 从 Kaggle 数据集中进行猫/狗检测。我启用了 cuda，看起来训练非常快。但是，数据加载（或处理）似乎非常缓慢。有什么方法可以加快速度吗？我试图玩女巫批量大小，但没有提供太多帮助。我还将 num_workers 从 0 更改为一些正数。从 0 到 2 可能会减少 1/3 的加载时间，增加更多不会产生额外的影响。还有其他方法可以加快加载速度吗？

这是我的粗略代码（我专注于学习，所以不是很有条理）：

import matplotlib.pyplot as plt

import torch
from torch import nn
from torch import optim
import torch.nn.functional as F
from torchvision import datasets, transforms, models

data_dir = 'Cat_Dog_data'

train_transforms = transforms.Compose([transforms.RandomRotation(30),
                                       transforms.RandomResizedCrop(224),
                                       transforms.RandomHorizontalFlip(),
                                       transforms.ToTensor(),
                                       transforms.Normalize([0.5, 0.5, 0.5],
                                                            [0.5, 0.5, 0.5])])
test_transforms = transforms.Compose([transforms.Resize(255),
                                      transforms.CenterCrop(224),
                                      transforms.ToTensor()])

# Pass transforms in here, then run the next cell to see how the transforms look
train_data = datasets.ImageFolder(data_dir + '/train',
                                  transform=train_transforms)
test_data = datasets.ImageFolder(data_dir + '/test', transform=test_transforms)

trainloader = torch.utils.data.DataLoader(train_data, batch_size=64,
                                          num_workers=16, shuffle=True,
                                          pin_memory=True)
testloader = torch.utils.data.DataLoader(test_data, batch_size=64,
                                         num_workers=16)

model = models.densenet121(pretrained=True)

# Freeze parameters so we don't backprop through them
for param in model.parameters():
    param.requires_grad = False

from collections import OrderedDict

classifier = nn.Sequential(OrderedDict([
    ('fc1', nn.Linear(1024, 500)),
    ('relu', nn.ReLU()),
    ('fc2', nn.Linear(500, 2)),
    ('output', nn.LogSoftmax(dim=1))
]))

model.classifier = classifier
model.cuda()
criterion = nn.NLLLoss()
optimizer = optim.Adam(model.parameters(), lr=0.003)

epochs = 30
steps = 0

import time

device = torch.device('cuda:0')

train_losses, test_losses = [], []
for e in range(epochs):
    running_loss = 0
    count = 0
    total_start = time.time()
    for images, labels in trainloader:
        start = time.time()
        images = images.cuda()
        labels = labels.cuda()

        optimizer.zero_grad()

        log_ps = model(images)
        loss = criterion(log_ps, labels)
        loss.backward()
        optimizer.step()
        elapsed = time.time() - start

        if count % 20 == 0:
            print("Optimized elapsed: ", elapsed, "count:", count)
            print("Total elapsed ", time.time() - total_start)
            total_start = time.time()
        count += 1

        running_loss += loss.item()
    else:
        test_loss = 0
        accuracy = 0
        for images, labels in testloader:
            images = images.cuda()
            labels = labels.cuda()
            with torch.no_grad():
                model.eval()
                log_ps = model(images)
                test_loss += criterion(log_ps, labels)
                ps = torch.exp(log_ps)
                top_p, top_class = ps.topk(1, dim=1)
                compare = top_class == labels.view(*top_class.shape)
                accuracy += compare.type(torch.FloatTensor).mean()
        model.train()
        train_losses.append(running_loss / len(trainloader))
        test_losses.append(test_loss / len(testloader))

        print("Epoch: /.. ".format(e + 1, epochs),
              "Training Loss: :.3f.. ".format(
                  running_loss / len(trainloader)),
              "Test Loss: :.3f.. ".format(test_loss / len(testloader)),
              "Test Accuracy: :.3f".format(accuracy / len(testloader)))

【参考方案1】：

torchvision 0.8.0 或更高版本

实际上，torchvision 现在在转换方面支持批处理和 GPU（这是在 torch.Tensors 上完成的，而不是 PIL 图像），因此应该将其用作初始改进。

有关此版本的更多信息，请参阅here。也可以用作torch.nn.Module，因此可以在模型内部使用，例如：

transforms = torch.nn.Sequential(
    T.RandomCrop(224),
    T.RandomHorizontalFlip(p=0.3),
    T.ConvertImageDtype(torch.float),
    T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
)

此外，这些操作可以进行 JIT 处理，可能会进一步提高性能。

torchvision 0.8.0（原答案）

增加batch_size 将无济于事，因为torchvision 在从磁盘加载单个图像时对其执行转换。

随着难度的增加，有几种方法可以加快数据加载速度：

缩短图片加载时间 在 RAM（或磁盘）中加载和规范化图像和缓存 生成转换并将其保存到磁盘 以批处理方式应用不可缓存的变换（旋转、翻转、裁剪） 预取

1。改善图片加载

通过安装Pillow-SIMD 而不是原来的pillow，可以获得简单的改进。它是一种直接替代品，可能会更快（或者至少对于您正在使用的Resize 声称如此）。

或者，您可以使用OpenCV 创建自己的数据加载和处理，因为有人说它更快或检查albumentations（虽然不能告诉您这些是否会提高性能并且可能会浪费大量时间除了学习经验没有收获）。

2。加载和规范化图像和缓存

您可以使用 Python 的 LRU Cache 功能来缓存一些输出。

您还可以使用torchdata，它的作用几乎与 PyTorch 的 torch.utils.data.Dataset 完全相同，但允许在 torchdata.Dataset 上使用简单的 cache() 缓存到磁盘或 RAM（或混合模式）（请参阅 github repository，免责声明：我是作者）。

记住：您必须加载和规范化图像、缓存，然后使用RandomRotation、RandomResizedCrop 和RandomHorizontalFlip（因为它们每次运行都会改变）。

3。生成转换并将它们保存到磁盘

您必须对图像执行大量转换，将它们保存到磁盘，然后再使用这个增强的数据集。再一次，这可以通过torchdata 完成，但在 I/O 和硬盘驱动器以及非常不雅的解决方案方面确实很浪费。此外，它是“静态的”，因此数据只会持续 X 个时期，它不会是带有增强功能的“无限”生成器。

4。批量转换

torchvision 不支持它，因此您必须自己编写这些函数。有关理由，请参阅this issue。 AFAIK 也没有其他第三方提供。对于大批量，它应该加快速度，但我认为实施是一个悬而未决的问题（如果我错了，请纠正我）。

5。预取

IMO 将是最难实施的（尽管考虑这个项目确实是一个好主意）。基本上，您在模型训练时为下一次迭代加载数据。 torch.utils.data.DataLoader 确实提供了，尽管存在一些问题（例如工作人员在加载数据后暂停）。您可以阅读 PyTorch thread 的相关信息（不确定，因为我没有自己验证）。此外，还有很多有价值的见解 provided by this comment 和 this blog post（虽然不确定它们的最新情况）。

总而言之，要显着改善数据加载，你需要弄脏你的手（或者也许有一些库在为 PyTorch 做这件事，如果是这样，我很想知道他们）。

还记得配置您的更改，请参阅torch.nn.bottleneck

编辑： DALI 项目可能值得一试，尽管 AFAIK 存在 RAM 内存随 epoch 数线性增长的问题。