Pytorch：使用data_prefetcher提升数据读取速度【自动将训练数据移动到GPU】

直接给出代码：

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 加速数据读取