火炬。 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

但我期待这样的结果，因为我在 Dataloader.

中指定了标志 pin_memory=True

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, but the gist of it is that this flag allows the x.cuda() operation (which you still have to execute as usually) to avoid one implicit CPU-to-CPU copy, which makes it a bit more performant. Additionally, with pinned memory tensors you can use 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 - 我不确定这是不是真的，请不要引用我的话那。

编辑：我相信您的基准测试没有抓住要点。它存在三个问题

您没有在 .cuda() 通话中使用 non_blocking=True。
您没有在 DataLoader 中使用多处理，这意味着大部分工作无论如何都是在主线程上同步完成的，超过了内存传输成本。
您没有在数据加载循环中执行任何 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.48 秒，没有内存固定的平均时间为 5.72 秒。

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

Pytorch. How does pin_memory work in Dataloader?

deep-learning

torch

pytorch