PyTorch：加快数据加载速度

时间：2020-04-23 17:40:16

标签： pytorch dataloader

我正在使用densitynet121从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 个答案:

答案 0 :(得分：10)

torchvision 0.8.0版本或更高版本

实际上，torchvision现在支持批量和GPU转换（这是在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])
)

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

torchvision <0.8.0（原始答案）

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

有两种方法可以提高难度，从而加快数据加载速度：

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

1。改善图像加载

通过安装Pillow-SIMD而不是原始的pillow可以轻松改进。它是一种替代产品，可能会更快（或者至少对于您使用的Resize而言如此）。

或者，您可以使用OpenCV创建自己的数据加载和处理，因为有人说它更快，或者检查albumentations（尽管无法告诉您它们是否可以提高性能，并且可能会很多的时间浪费，除了学习经验之外没有任何收获。

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

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

您还可以使用torchdata，其行为几乎与PyTorch的{​​{1}}相同，但是可以通过在torch.utils.data.Dataset上使用简单的cache()来缓存到磁盘或RAM（或混合模式）中（参见github repository，免责声明：我是作者）。

请记住：您必须加载并规范化图像，缓存，然后再使用torchdata.Dataset，RandomRotation和RandomResizedCrop（因为它们每次都更改运行）。

3。产生转换并将其保存到磁盘

您必须对图像执行很多转换，将其保存到磁盘上，然后再使用此增强的数据集。再次可以使用RandomHorizontalFlip来完成，但在I / O和硬盘驱动器以及非常笨拙的解决方案方面确实非常浪费。此外，它是“静态”的，因此数据只会持续X个纪元，而不会是具有扩充功能的“无限”生成器。

4。批量转换

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

5。预取

IMO将是最难实施的（尽管考虑该项目的想法非常好）。基本上，您在模型训练时为下一次迭代加载数据。 torchvision 是否提供，尽管存在一些问题（例如，工作人员在加载数据后暂停）。您可以阅读PyTorch thread的相关内容（不确定该内容，因为我没有自己进行验证）。此外，provided by this comment和this blog post还提供了许多有价值的见解（尽管不确定这些消息的最新程度）。

总而言之，要大大改善数据加载量，您需要变得非常肮脏（或者也许有一些图书馆为PyTorch做了一些此类工作，如果是的话，我很想知道他们）。

还请记住分析您的更改，请参阅torch.nn.bottleneck

编辑： DALI项目可能值得一试，尽管AFAIK的RAM内存随着时代的增长呈线性增长，但存在一些问题。

相关问题
最新问题