小编Iva*_*gov的帖子

我想了解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()) …