From c0d180e72fa23baffa76d199b0ea0cc74acd766f Mon Sep 17 00:00:00 2001 From: Hunor Laczko Date: Thu, 6 Apr 2023 17:24:27 +0000 Subject: [PATCH] add infinite dataloader --- utils/utils.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/utils/utils.py b/utils/utils.py index 337bc40..d1d247a 100644 --- a/utils/utils.py +++ b/utils/utils.py @@ -1,6 +1,7 @@ import os import sys import numpy as np +from torch.utils.data import DataLoader @@ -25,3 +26,20 @@ class Accumulator(): self.dict = dict() +class InfiniteDataLoader(DataLoader): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # Initialize an iterator over the dataset. + self.dataset_iterator = super().__iter__() + + def __iter__(self): + return self + + def __next__(self): + try: + batch = next(self.dataset_iterator) + except StopIteration: + # Dataset exhausted, use a new fresh iterator. + self.dataset_iterator = super().__iter__() + batch = next(self.dataset_iterator) + return batch \ No newline at end of file