火炬。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内存。这里提供了好处和理由，但要点是，此标志允许x.cuda()操作（您仍然必须照常执行）避免一个隐式的CPU到CPU复制，这使它性能更高。此外，通过固定的内存张量，您可以用于x.cuda(non_blocking=True)相对于主机异步执行复制。在某些情况下，这可能会导致性能提高，即如果您的代码结构如下

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

由于开始的复制1.是异步的，因此2.在复制过程中它不会阻止继续进行，因此两者可以并排发生（这就是收益）。由于步骤3.必须x已经复制到GPU，因此无法执行，直到1.完成为止-因此只能执行1.并且2.可以重叠，并且3.肯定会在之后进行。2.因此，的持续时间是您可以预期节省的最长时间non_blocking=True。如果没有non_blocking=True您的CPU，则在继续之前将等待空闲以完成传输2.。

注意：步骤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。