如何解决正确使用PyTorch的non_blocking = True进行数据预取
我正在研究在GPU上训练模型时将数据从CPU预取到GPU中的情况。通过GPU模型训练将CPU到GPU的数据传输重叠在一起似乎需要两者
- 使用
data = data.cuda(non_blocking=True)
将数据传输到GPU - 使用
train_loader = DataLoader(...,pin_memory=True)
将数据固定到CPU内存
但是,我无法理解在此official PyTorch example中如何执行无阻塞传输,特别是以下代码块:
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执行数据预取的正确方法是什么?
解决方法
要在PyTorch上正确实现GPU预取,您必须将for循环转换为while循环。
应使用iter
函数将DataLoader更改为迭代器,例如iterator = iter(loader)
。
在while循环内的每一步使用next(iterator)
来获取下一个迷你批处理。
可以通过从迭代器中捕获StopIteration
来检测DataLoader的结尾。
在引发StopIteration
时,使用标志结束while循环。
版权声明:本文内容由互联网用户自发贡献,该文观点与技术仅代表作者本人。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如发现本站有涉嫌侵权/违法违规的内容, 请发送邮件至 dio@foxmail.com 举报,一经查实,本站将立刻删除。