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

Posted 2023-03-12

【中文标题】火炬。在 Dataloader 中 pin_memory 是如何工作的？

【英文标题】：Pytorch. How does pin_memory work in Dataloader?

【发布时间】：2019-08-29 00:13:45

【问题描述】：

我想了解 Dataloader 中 pin_memory 的工作原理。

根据文档：

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 只会让事情变慢。 有人可以解释一下这种行为吗？

【问题讨论】：

我已编辑我的答案以响应您的基准测试。下次请发表评论，因为我只是偶然发现您的问题已更新

【参考方案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 - 我不确定这是否属实，请不要引用我的话。

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

您没有在 .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.48s 内存固定和 5.72s 没有。

【讨论】：

这是否意味着额外的 RAM 使用？我们什么时候不应该使用它？谢谢

我不知道技术细节和确切的后果。我认为没有使用任何额外的 RAM，但由于无法调出它，因此操作系统可能无法调出您的程序和 OOM，而它通常可以从中恢复。

你知道.to(non_blocking=True)在pin_memory==False时的预期行为吗？

我不明白的是，如果 .cuda 操作与 CPU 操作同时出现。怎么保证发给.cuda的x是处理后的x，而不是原来的？

这不是额外的内存使用，而是操作系统无法移动的一块内存，如果内存不足则交换到磁盘等。因此它使操作系统的工作更加困难，并且您可以固定多少内存是有限制的。