PyTorch：加快数据加载

Question

我正在使用 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)))

Answer 1

torchvision0.8.0版本或更高版本

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

有关此版本的更多信息，请参见此处。那些也充当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有人说它更快或检查专辑（尽管无法告诉您这些是否会提高性能，并且可能会浪费大量时间，除了学习经验之外没有任何收获）。

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

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

您还可以使用torchdatawhich 的行为几乎与 PyTorch 的完全一样，torch.utils.data.Dataset但允许缓存到磁盘或 RAM（或混合模式），只需简单cache()打开torchdata.Dataset（请参阅github 存储库，免责声明：我是作者）。

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

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

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

4. 批量转换

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

5. 预取

IMO 将是最难实施的（尽管考虑该项目的一个非常好的主意）。基本上，您在模型训练时为下一次迭代加载数据。torch.utils.data.DataLoader 确实提供了它，尽管存在一些问题（例如工作人员在数据加载后暂停）。您可以阅读有关它的PyTorch 线程（不确定，因为我没有亲自验证）。此外，此评论和此博客文章提供了许多有价值的见解（尽管不确定它们的最新情况）。

总而言之，要显着改善数据加载，您需要让自己的手变得非常脏（或者可能有一些库为 PyTorch 执行了这些操作，如果是这样，我很想了解它们）。

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

编辑： DALI项目可能值得一试，尽管 AFAIK 它在 RAM 内存随时期数线性增长方面存在一些问题。