A general prefetch generator. Ref: https://stackoverflow.com/questions/7323664/python-generator-pre-fetch Args: generator: Python generator. num_prefetch_queue (int): Number of prefetch queue.
| 5 | |
| 6 | |
| 7 | class PrefetchGenerator(threading.Thread): |
| 8 | """A general prefetch generator. |
| 9 | |
| 10 | Ref: |
| 11 | https://stackoverflow.com/questions/7323664/python-generator-pre-fetch |
| 12 | |
| 13 | Args: |
| 14 | generator: Python generator. |
| 15 | num_prefetch_queue (int): Number of prefetch queue. |
| 16 | """ |
| 17 | |
| 18 | def __init__(self, generator, num_prefetch_queue): |
| 19 | threading.Thread.__init__(self) |
| 20 | self.queue = Queue.Queue(num_prefetch_queue) |
| 21 | self.generator = generator |
| 22 | self.daemon = True |
| 23 | self.start() |
| 24 | |
| 25 | def run(self): |
| 26 | for item in self.generator: |
| 27 | self.queue.put(item) |
| 28 | self.queue.put(None) |
| 29 | |
| 30 | def __next__(self): |
| 31 | next_item = self.queue.get() |
| 32 | if next_item is None: |
| 33 | raise StopIteration |
| 34 | return next_item |
| 35 | |
| 36 | def __iter__(self): |
| 37 | return self |
| 38 | |
| 39 | |
| 40 | class PrefetchDataLoader(DataLoader): |