Prefetch version of dataloader. Ref: https://github.com/IgorSusmelj/pytorch-styleguide/issues/5# TODO: Need to test on single gpu and ddp (multi-gpu). There is a known issue in ddp. Args: num_prefetch_queue (int): Number of prefetch queue. kwargs (dict): Ot
| 38 | |
| 39 | |
| 40 | class PrefetchDataLoader(DataLoader): |
| 41 | """Prefetch version of dataloader. |
| 42 | |
| 43 | Ref: |
| 44 | https://github.com/IgorSusmelj/pytorch-styleguide/issues/5# |
| 45 | |
| 46 | TODO: |
| 47 | Need to test on single gpu and ddp (multi-gpu). There is a known issue in |
| 48 | ddp. |
| 49 | |
| 50 | Args: |
| 51 | num_prefetch_queue (int): Number of prefetch queue. |
| 52 | kwargs (dict): Other arguments for dataloader. |
| 53 | """ |
| 54 | |
| 55 | def __init__(self, num_prefetch_queue, **kwargs): |
| 56 | self.num_prefetch_queue = num_prefetch_queue |
| 57 | super(PrefetchDataLoader, self).__init__(**kwargs) |
| 58 | |
| 59 | def __iter__(self): |
| 60 | return PrefetchGenerator(super().__iter__(), self.num_prefetch_queue) |
| 61 | |
| 62 | |
| 63 | class CPUPrefetcher(): |