|
|
|
|
|
|
|
|
|
import torch |
|
|
|
|
|
class DataPrefetcher: |
|
""" |
|
DataPrefetcher is inspired by code of following file: |
|
https://github.com/NVIDIA/apex/blob/master/examples/imagenet/main_amp.py |
|
It could speedup your pytorch dataloader. For more information, please check |
|
https://github.com/NVIDIA/apex/issues/304#issuecomment-493562789. |
|
""" |
|
|
|
def __init__(self, loader): |
|
self.loader = iter(loader) |
|
self.stream = torch.cuda.Stream() |
|
self.input_cuda = self._input_cuda_for_image |
|
self.record_stream = DataPrefetcher._record_stream_for_image |
|
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.input_cuda() |
|
self.next_target = self.next_target.cuda(non_blocking=True) |
|
|
|
def next(self): |
|
torch.cuda.current_stream().wait_stream(self.stream) |
|
input = self.next_input |
|
target = self.next_target |
|
if input is not None: |
|
self.record_stream(input) |
|
if target is not None: |
|
target.record_stream(torch.cuda.current_stream()) |
|
self.preload() |
|
return input, target |
|
|
|
def _input_cuda_for_image(self): |
|
self.next_input = self.next_input.cuda(non_blocking=True) |
|
|
|
@staticmethod |
|
def _record_stream_for_image(input): |
|
input.record_stream(torch.cuda.current_stream()) |
|
|
