我正在研究在GPU上训练模型时将数据从CPU预取到GPU中。与 GPU 模型训练重叠的 CPU 到 GPU 数据传输似乎需要两者
- 使用
data = data.cuda(non_blocking=True)将数据传输到 GPU - 使用
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)是阻塞的，因此在计算模型输出之前，for循环的下一个i迭代中的images.cuda()和target.cuda()不会发生，这意味着在下一个循环迭代中不会进行预取。
如果这是正确的，那么将数据预取到 GPU 的正确方法是什么？
我认为你关闭的地方是output = model(images)是一个同步点。似乎计算是由 GPU 的不同部分处理的。引用官方 PyTorch 文档：
此外，固定张量或存储后，可以使用异步 GPU 副本。只需将一个额外的
non_blocking=True参数传递给to()或cuda()电话。这可用于重叠数据 使用计算进行传输。
要在 PyTorch 上正确实现 GPU 预取，必须将 for 循环转换为 while-loop。
数据加载器应使用iter函数更改为迭代器，例如iterator = iter(loader).
在 while-循环中的每个步骤中使用next(iterator)来获取下一个小批量。
可以通过从迭代器中捕获StopIteration来检测数据加载器的结束。
使用标志在引发StopIteration时结束 while 循环。