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

Question

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

1. 将数据传输到 GPU 使用 data = data.cuda(non_blocking=True)
2. 将数据固定到 CPU 内存使用 train_loader = DataLoader(..., pin_memory=True)

但是，我无法理解在这个官方 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只需向 a to()或调用传递一个附加参数即可cuda()。这可用于将数据传输与计算重叠。

Answer 2

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

应使用函数将 DataLoader 更改为迭代器iter，例如iterator = iter(loader)。

next(iterator)在 while 循环内的每个步骤中使用以获得下一个小批量。

StopIteration可以通过从迭代器捕获 来检测 DataLoader 的结束。

StopIteration当引发时，使用标志来结束 while 循环。