Pytorch。 pin_memory如何在Dataloader中运行？

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)相对于主机异步执行复制。这可以在某些情况下提高性能，即如果您的代码结构为

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

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

注意：也许步骤2.也可以包含GPU操作，只要它们不需要x - 我不确定这是否属实，请不要引用我。

编辑：我相信你错过了你的基准点。它有三个问题

1. 你没有在non_blocking=True电话中使用.cuda()。
2. 你没有在你的DataLoader中使用多处理，这意味着大多数工作都是在主线程上同步完成的，超过了内存传输成本。
3. 您没有在数据加载循环中执行任何CPU工作（除了.cuda()调用），因此没有工作可以覆盖内存传输。

更接近于如何使用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。