直接给出代码:
class DataPreFetcher(object):
def __init__(self, loader):
self.loader = iter(loader)
self.stream = torch.cuda.Stream()
self.preload()
def preload(self):
try:
self.next_data = next(self.loader)
except StopIteration:
self.next_data = None
return
with torch.cuda.stream(self.stream):
for k, v in self.next_data.items():
if isinstance(v, torch.Tensor):
self.next_data[k] = self.next_data[k].cuda(non_blocking=True)
def next(self):
torch.cuda.current_stream().wait_stream(self.stream)
data = self.next_data
self.preload()
return data
class data_prefetcher():
def __init__(self, loader):
#loader 1:real
#loader 2:fake
self.stream = torch.cuda.Stream()
self.loader = iter(loader)
self.preload()
def preload(self):
try:
self.next_input, self.next_target = next(self.loader)
except StopIteration:
self.next_input = None
self.next_target = None
return
with torch.cuda.stream(self.stream):
self.next_input = self.next_input.cuda(non_blocking=True).float()
self.next_target = self.next_target.cuda(non_blocking=True).long()
def next(self):
torch.cuda.current_stream().wait_stream(self.stream)
input = self.next_input
target = self.next_target
self.preload()
return input, target
class DataPreFetcher(object):
def __init__(self, dataLoader):
self.stream = torch.cuda.Stream()
self.dataLoader = iter(dataLoader)
self.preload()
def preload(self):
try:
self.next_batch_data = next(self.dataLoader)
except StopIteration:
self.next_batch_data = None
return
with torch.cuda.stream(self.stream):
for k, v in self.next_batch_data.items():
if isinstance(v, torch.Tensor):
self.next_batch_data[k] = self.next_batch_data[k].cuda(non_blocking=True)
def next(self):
torch.cuda.current_stream().wait_stream(self.stream)
batch_data = self.next_batch_data
self.preload()
return batch_data
参考资料:
pytorch使用data_prefetcher提升数据读取速度
使用DataPrefetcher加速PyTorch的dataloader
【pytorch】给训练踩踩油门-- Pytorch 加速数据读取