Question

我想了解pin_memory在Dataloader中的工作方式。

根据文档。

pin_memory (bool, optional) – If True, the data loader will copy tensors into CUDA pinned memory before returning them.

下面是自包含的代码示例。

import torchvision
import torch

print('torch.cuda.is_available()', torch.cuda.is_available())
train_dataset = torchvision.datasets.CIFAR10(root='cifar10_pytorch', download=True, transform=torchvision.transforms.ToTensor())
train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=64, pin_memory=True)
x, y = next(iter(train_dataloader))
print('x.device', x.device)
print('y.device', y.device)

我得到了输出。

torch.cuda.is_available() True
x.device cpu
y.device cpu

但是我期待这样的事情，因为我在pin_memory=True中指定了标志Dataloader。

torch.cuda.is_available() True
x.device cuda:0
y.device cuda:0

我还运行一些基准测试：

import torchvision
import torch
import time
import numpy as np

pin_memory=True
train_dataset =torchvision.datasets.CIFAR10(root='cifar10_pytorch', download=True, transform=torchvision.transforms.ToTensor())
train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=64, pin_memory=pin_memory)
print('pin_memory:', pin_memory)
times = []
n_runs = 10
for i in range(n_runs):
    st = time.time()
    for bx, by in train_dataloader:
        bx, by = bx.cuda(), by.cuda()
    times.append(time.time() - st)
print('average time:', np.mean(times))

我得到了以下结果。

pin_memory: False
average time: 6.5701503753662

pin_memory: True
average time: 7.0254474401474

所以pin_memory=True只会使事情变慢。有人可以向我解释这种行为吗？

Answer 1

考虑到所使用的术语相当小众，该文档可能过于简洁。用CUDA术语来说，固定内存并不意味着GPU内存，而是非页面CPU内存。提供了好处和理由here，但要点在于，此标志允许x.cuda()操作（您仍然必须照常执行）以避免一个隐式的CPU到CPU复制，这使其性能更高。此外，使用固定的内存张量，您可以使用x.cuda(non_blocking=True)相对于主机异步执行复制。在某些情况下，这可能会导致性能提高，即如果您的代码结构为

x.cuda(non_blocking=True)
执行一些CPU操作
使用x执行GPU操作。

由于在1.中发起的复制是异步的，因此它不会阻止2.在复制过程中继续进行，因此两者可以并排发生（这是收益）。由于步骤3.要求x已被复制到GPU，因此无法执行，直到1.完成-因此只有1.和2.可以重叠，然后3.肯定会发生。因此，2.的持续时间是可以预期用non_blocking=True保存的最长时间。如果没有non_blocking=True，您的CPU将等待空闲，等待传输完成，然后再继续进行2.。

注意：步骤2.可能还包含GPU操作，只要它们不需要x-我不确定这是否成立，请不要在此引用我的意思。

编辑：我相信您在基准测试中遗漏了重点。有三个问题

您在non_blocking=True通话中没有使用.cuda()。
您没有在DataLoader中使用多处理，这意味着无论如何大多数工作都是在主线程上同步完成的，从而大大节省了内存传输成本。
除了在.cuda()调用之外，您没有在数据加载循环中执行任何CPU工作，因此没有工作要覆盖内存传输。

一个更接近pin_memory用法的基准将是

import torchvision, torch, time
import numpy as np

pin_memory = True
batch_size = 1024 # bigger memory transfers to make their cost more noticable
n_workers = 6 # parallel workers to free up the main thread and reduce data decoding overhead
train_dataset =torchvision.datasets.CIFAR10(
    root='cifar10_pytorch',
    download=True,
    transform=torchvision.transforms.ToTensor()
)   
train_dataloader = torch.utils.data.DataLoader(
    train_dataset,
    batch_size=batch_size,
    pin_memory=pin_memory,
    num_workers=n_workers
)   
print('pin_memory:', pin_memory)
times = []
n_runs = 10

def work():
    # emulates the CPU work done
    time.sleep(0.1)

for i in range(n_runs):
    st = time.time()
    for bx, by in train_dataloader:
       bx, by = bx.cuda(non_blocking=pin_memory), by.cuda(non_blocking=pin_memory)
       work()
   times.append(time.time() - st)
print('average time:', np.mean(times))

对于具有内存固定功能的计算机，平均为5.48s，而对于不带内存固定的计算机，平均为5.72s。

火炬。 pin_memory如何在Dataloader中工作？

1 个答案: