gpu training util, to accelerate dataloader's speed

async_queue.py

from queue import Queue
from threading import Thread


class AsyncQueue:
    def __init__(self, generator, cache_size=10):
        self._q = Queue(maxsize=cache_size)
        self._generator = generator
        self._input_thread = Thread(target=self.append_to_queue)
        self._input_thread.setDaemon(True)
        self._input_thread.start()

    def append_to_queue(self):
        for item in self._generator:
            self._q.put(item, block=True)

    def __iter__(self):
        return self

    def __next__(self):
        return self._q.get(block=True)