How to load/fetch the next data batches for the next epoch, during the current epoch?

user_2838606 0
Viewed 125

I know that since PyTorch 1.7.0 it is possible to prefetch some batches before an epoch begins. However, this does not make it possible to fetch batches while the operations within an epoch are being performed and before the next epoch begins. Based on this thread, it seems that it should be possible to use a Sampler to load batches during an epoch, and before the next epoch begins. However, I cannot wrap my head around how I can use a Sampler to achieve this.

Can anyone provide a code sample for a Sampler that allows fetching samples during an epoch?

pytorch computer-vision machine-learning python-3.x
1 Answers

You can prefetch the next batches from iterator in a background thread.

class _ThreadedIterator(threading.Thread):
    """
    Prefetch the next queue_length items from iterator in a background thread.

    Example:
    >> for i in bg_iterator(range(10)):
    >>     print(i)
    """

    class _End:
        pass

    def __init__(self, generator: Iterable, maxsize: int) -> None:
        threading.Thread.__init__(self)
        self.queue: Queue = Queue(maxsize)
        self.generator = generator
        self.daemon = True
        self.start()

    def run(self) -> None:
        for item in self.generator:
            self.queue.put(item)
        self.queue.put(self._End)

    def __iter__(self) -> Any:
        return self

    def __next__(self) -> Any:
        next_item = self.queue.get()
        if next_item == self._End:
            raise StopIteration
        return next_item

    # Required for Python 2.7 compatibility
    def next(self) -> Any:
        return self.__next__()


def bg_iterator(iterable: Iterable, maxsize: int) -> Any:
    return _ThreadedIterator(iterable, maxsize=maxsize)

UPD. Usage:

model = model.to(device, non_blocking=True)
for inputs, targets in bg_iterator(data_loader, maxsize=2):
    inputs = inputs.to(device, non_blocking=True)
    targets = targets.to(device, non_blocking=True)

example

user_39600380
Related