正确使用 PyTorch 的 non_blocking=True 进行数据预取

Question

我正在研究当模型在 GPU 上训练时将数据从 CPU 预取到 GPU 中。将 CPU 到 GPU 的数据传输与 GPU 模型训练重叠似乎两者都需要

1. 使用

   data = data.cuda(non_blocking=True)

将数据传输到 GPU
2. 使用

   train_loader = DataLoader(..., pin_memory=True)

将数据固定到 CPU 内存

但是，我无法理解在这个官方 PyTorch 示例中如何执行非阻塞传输，特别是这个代码块：

for i, (images, target) in enumerate(train_loader):
        # measure data loading time
        data_time.update(time.time() - end)

        if args.gpu is not None:
            images = images.cuda(args.gpu, non_blocking=True)
        if torch.cuda.is_available():
            target = target.cuda(args.gpu, non_blocking=True)

        # compute output
        output = model(images)
        loss = criterion(output, target)

images.cuda(non_blocking=True)

和

target.cuda(non_blocking=True)

不必在执行

output = model(images)

之前完成。由于这是一个同步点，

images

必须首先完全传输到CUDA设备，因此数据传输步骤实际上不再是非阻塞的。

由于

output = model(images)

是阻塞的，因此在计算模型输出之前，

images.cuda()

循环的下一个

target.cuda()

迭代中的

i

和

for

不会发生，这意味着在下一个循环迭代中不会预取。

如果这是正确的，那么向 GPU 执行数据预取的正确方法是什么？

Answer 1

我认为你的出发点是

output = model(images)

是一个同步点。看起来计算是由 GPU 的不同部分处理的。引用自官方 PyTorch 文档：

此外，一旦固定张量或存储，您就可以使用异步 GPU 副本。只需将额外的

non_blocking=True

参数传递给

to()

或拨打

cuda()

电话。 这可用于重叠数据 通过计算进行传输.

Answer 2

要在 PyTorch 上正确实现 GPU 预取，必须将 for 循环转换为 while 循环。

DataLoader 应使用

iter

函数更改为迭代器，例如

iterator = iter(loader)

。

在 while 循环内的每一步使用

next(iterator)

来获取下一个小批量。

可以通过从迭代器捕获

StopIteration

来检测 DataLoader 的结束。

当

StopIteration

升高时，使用标志结束 while 循环。