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.
| 11 | |
| 12 | |
| 13 | class DataPrefetcher: |
| 14 | """ |
| 15 | DataPrefetcher is inspired by code of following file: |
| 16 | https://github.com/NVIDIA/apex/blob/master/examples/imagenet/main_amp.py |
| 17 | It could speedup your pytorch dataloader. For more information, please check |
| 18 | https://github.com/NVIDIA/apex/issues/304#issuecomment-493562789. |
| 19 | """ |
| 20 | |
| 21 | def __init__(self, loader): |
| 22 | self.loader = iter(loader) |
| 23 | self.stream = torch.cuda.Stream() |
| 24 | self.input_cuda = self._input_cuda_for_image |
| 25 | self.record_stream = DataPrefetcher._record_stream_for_image |
| 26 | self.preload() |
| 27 | |
| 28 | def preload(self): |
| 29 | try: |
| 30 | self.next_input, self.next_target, _, _ = next(self.loader) |
| 31 | except StopIteration: |
| 32 | self.next_input = None |
| 33 | self.next_target = None |
| 34 | return |
| 35 | |
| 36 | with torch.cuda.stream(self.stream): |
| 37 | self.input_cuda() |
| 38 | self.next_target = self.next_target.cuda(non_blocking=True) |
| 39 | |
| 40 | def next(self): |
| 41 | torch.cuda.current_stream().wait_stream(self.stream) |
| 42 | input = self.next_input |
| 43 | target = self.next_target |
| 44 | if input is not None: |
| 45 | self.record_stream(input) |
| 46 | if target is not None: |
| 47 | target.record_stream(torch.cuda.current_stream()) |
| 48 | self.preload() |
| 49 | return input, target |
| 50 | |
| 51 | def _input_cuda_for_image(self): |
| 52 | self.next_input = self.next_input.cuda(non_blocking=True) |
| 53 | |
| 54 | @staticmethod |
| 55 | def _record_stream_for_image(input): |
| 56 | input.record_stream(torch.cuda.current_stream()) |
| 57 | |
| 58 | |
| 59 | def random_resize(data_loader, exp, epoch, rank, is_distributed): |