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

曾经终败给现在 2023-09-27 09:56 102阅读 0赞

直接给出代码：

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提升数据读取速度][pytorch_data_prefetcher]  
[使用DataPrefetcher加速PyTorch的dataloader][DataPrefetcher_PyTorch_dataloader]  
[【pytorch】给训练踩踩油门-- Pytorch 加速数据读取][pytorch_-- Pytorch]

[pytorch_data_prefetcher]: https://blog.csdn.net/Mr_health/article/details/122837461
[DataPrefetcher_PyTorch_dataloader]: https://www.jianshu.com/p/21bce3a05070
[pytorch_-- Pytorch]: https://blog.csdn.net/shwan_ma/article/details/103331166