PyTorch：加速数据加载
PyTorch: Speed up data loading
我正在使用 densenet121 从 Kaggle 数据集进行 cat/dog 检测。我启用了 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)))
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 as some say it's faster or check albumentations 创建您自己的数据加载和处理（尽管无法告诉您这些是否会提高性能，并且可能会浪费大量时间，除了学习经验之外毫无收获） .
2。加载和规范化图像和缓存
您可以使用 Python 的 LRU Cache 功能来缓存一些输出。
您还可以使用 torchdata，它的行为几乎与 PyTorch 的 torch.utils.data.Dataset 完全相同，但允许在 [=24= 上使用简单的 cache() 缓存到磁盘或 RAM（或混合模式） ]（参见 github repository、免责声明：我是作者）。
记住： 你必须加载和规范化图像、缓存，然后使用 RandomRotation、RandomResizedCrop 和 RandomHorizontalFlip（因为那些每次更改 运行).
3。生成转换并将它们保存到磁盘
您必须对图像执行大量 转换，将它们保存到磁盘，然后使用这个增强的数据集。再一次，这可以用 torchdata 来完成，但是当涉及到 I/O 和硬盘驱动器以及非常不雅的解决方案时，它真的很浪费。此外，它是“静态的”，因此数据只会持续 X 个纪元，它不会是带有扩充的“无限”生成器。
4。批量转换
torchvision 不支持它，因此您必须自己编写这些函数。参见 this issue 的理由。 AFAIK 也没有其他第 3 方提供它。对于大批量它应该加快速度，但我认为实施是一个悬而未决的问题（如果我错了请纠正我）。
5。预取
IMO 将是最难实施的（尽管对项目来说是一个非常好的想法来考虑它）。基本上，您在模型训练时为下一次迭代加载数据。 torch.utils.data.DataLoader 确实提供了它，尽管有一些问题（比如工作人员在加载数据后暂停）。您可以阅读 PyTorch thread about it (not sure about it as I didn't verify on my own). Also, a lot of valuable insight provided by this comment and this blog post（尽管不确定它们的最新情况）。
总而言之，要显着改善数据加载，您需要亲自动手（或者也许有一些库在为 PyTorch 执行此操作，如果是这样，我很想知道他们）。
还记得分析您的更改，请参阅 torch.nn.bottleneck
编辑： DALI 项目可能值得一试，尽管据我所知，它在 RAM 内存随时代数线性增长方面存在一些问题。
我正在使用 densenet121 从 Kaggle 数据集进行 cat/dog 检测。我启用了 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)))
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 as some say it's faster or check albumentations 创建您自己的数据加载和处理（尽管无法告诉您这些是否会提高性能，并且可能会浪费大量时间，除了学习经验之外毫无收获） .
2。加载和规范化图像和缓存
您可以使用 Python 的 LRU Cache 功能来缓存一些输出。
您还可以使用 torchdata，它的行为几乎与 PyTorch 的 torch.utils.data.Dataset 完全相同，但允许在 [=24= 上使用简单的 cache() 缓存到磁盘或 RAM（或混合模式） ]（参见 github repository、免责声明：我是作者）。
记住： 你必须加载和规范化图像、缓存，然后使用 RandomRotation、RandomResizedCrop 和 RandomHorizontalFlip（因为那些每次更改 运行).
3。生成转换并将它们保存到磁盘
您必须对图像执行大量 转换，将它们保存到磁盘，然后使用这个增强的数据集。再一次，这可以用 torchdata 来完成，但是当涉及到 I/O 和硬盘驱动器以及非常不雅的解决方案时，它真的很浪费。此外，它是“静态的”，因此数据只会持续 X 个纪元，它不会是带有扩充的“无限”生成器。
4。批量转换
torchvision 不支持它，因此您必须自己编写这些函数。参见 this issue 的理由。 AFAIK 也没有其他第 3 方提供它。对于大批量它应该加快速度，但我认为实施是一个悬而未决的问题（如果我错了请纠正我）。
5。预取
IMO 将是最难实施的（尽管对项目来说是一个非常好的想法来考虑它）。基本上，您在模型训练时为下一次迭代加载数据。 torch.utils.data.DataLoader 确实提供了它，尽管有一些问题（比如工作人员在加载数据后暂停）。您可以阅读 PyTorch thread about it (not sure about it as I didn't verify on my own). Also, a lot of valuable insight provided by this comment and this blog post（尽管不确定它们的最新情况）。
总而言之，要显着改善数据加载，您需要亲自动手（或者也许有一些库在为 PyTorch 执行此操作，如果是这样，我很想知道他们）。
还记得分析您的更改，请参阅 torch.nn.bottleneck
编辑： DALI 项目可能值得一试，尽管据我所知，它在 RAM 内存随时代数线性增长方面存在一些问题。