Deep learning 皮托克。在Dataloader中pin_内存是如何工作的？

我想了解Dataloader中的pin_内存是如何工作的

根据文件：

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

中指定了flag

pin\u 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

只会让事情变得更慢。

---

有人能给我解释一下这种行为吗？

鉴于所使用的术语相当合适，文档可能过于简洁。在CUDA术语中，固定内存并不意味着GPU内存，而是非分页CPU内存。虽然提供了好处和基本原理，但其要点是该标志允许

x.cuda（）

操作（您仍然必须像往常一样执行）避免一个隐式的CPU到CPU的复制，从而使其性能更高。此外，使用固定内存张量，您可以使用它来针对主机异步执行复制。在某些情况下，这可能会提高性能，即如果您的代码是按

x.cuda（非阻塞=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

您的

数据加载器中没有使用多处理，这意味着大部分工作无论如何都是在主线程上同步完成的，超过了内存传输成本

在数据加载循环中，您没有执行任何CPU工作（除了

.cuda（）

调用），因此内存传输不会覆盖任何工作

更接近于如何使用

pin_内存

的基准是

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秒的内存固定。

这是否意味着额外的RAM使用？我们什么时候不用它？谢谢，我不知道技术细节和确切后果。我认为没有使用任何额外的RAM，但由于它不能被调出，操作系统可能无法调出您的程序和OOM，而这种情况通常是可以恢复的。您知道

.to（non_blocking=True）

在

pin_memory==False

时的预期行为吗，如果

.cuda

操作与CPU操作同时出现。我们如何保证发送到

.cuda

的

x

是经过处理的

x

，而不是原始的？这不是额外的内存使用，而是操作系统无法移动的内存块，如果内存不足，则交换到磁盘，等等。因此，这使得操作系统的工作更加困难，你可以锁定多少内存是有限制的。我编辑了我的答案以响应你的基准测试。下次请留下评论，因为我只是碰巧注意到你的问题被更新了