直接给出代码：

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