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

torchvision 0.8.0版本或更高版本

实际上，当涉及到转换时，torchvision现在支持批处理和图形处理器(这是在torch.Tensor上完成的，而不是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])
)

此外，这些操作可以是JITed，可能会进一步提高性能。

torchvision < 0.8.0 (原始答案)

增加batch_size不会有什么帮助，因为torchvision会在从磁盘加载的单个图像上执行转换。

有几种方法可以在增加难度的情况下加速数据加载：

在批处理manner

• Prefetching

中应用不可缓存的变换(旋转、翻转、裁剪)

• 改善图像加载时间
• 在内存中(或在disk)
• Produce变换时)加载和标准化图像并将其保存到磁盘
• 应用不可缓存的变换(旋转、翻转、裁剪)

1.改善图片加载

通过安装Pillow-SIMD而不是原来的pillow，可以获得简单的改进。它是一个即用型的替代品，可能会更快(至少对于你正在使用的Resize是这样宣称的)。

或者，你可以使用OpenCV创建你自己的数据加载和处理，因为有些人说它更快，或者检查albumentations (虽然不能告诉你这些是否会提高性能，可能会浪费很多时间，除了学习经验之外，没有任何收获)。

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

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

您也可以使用torchdata，它的行为与PyTorch的torch.utils.data.Dataset几乎完全相同，但允许缓存到磁盘或在内存(或混合模式)中使用torchdata.Dataset上的简单cache() (参见github repository，免责声明: i'm author)。

记住：你必须加载和规格化图像，缓存，然后使用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项目可能值得一试，尽管它存在一些内存随着纪元数线性增长的问题。