Implementing an "infinite loop" Dataset & DataLoader combo

username: concretevitamin
cooked: <p>I’d like to implement an infinite loop Dataset &amp; DataLoader.  Here’s what I tried:</p>
<pre><code class="lang-auto">class Infinite(Dataset):
    def __len__(self):
        return HPARAMS.batch_size
#         return 1&lt;&lt;30 # This causes huge memory usage.
    def __getitem__(self, idx):
        """Randomly generates one new example."""
        return sample_func_to_be_parallelized()

infinite_loader = DataLoader(
    dataset=Infinite(), 
    batch_size=HPARAMS.batch_size, 
    num_workers=16,
    worker_init_fn=lambda worker_id: np.random.seed(worker_id),  
)

while True:
    for idx, data in enumerate(infinite_loader):
        # forward + backward on "data"
</code></pre>
<p>As you can see, the main challenge here is the <strong>len</strong>() method.  If I put a large enough number there, like 1&lt;&lt;30, the symptom is memory usage will JUMP TO 10+GB on the first iteration of train loop.  After a while the workers are killed presumably due to OOM.</p>
<p>If I put a small number there, like 1 or BATCH_SIZE, the sampled “data” in the train loop will be periodically duplicated.  This is not what I want as I’d like new data to be generated &amp; trained on at every iteration.</p>
<p>I’m guessing the culprit of the excessive memory usage is somewhere in the stack, a bunch of things are cached.  Upon a casual look at the Python side of things I can’t pinpoint where.</p>
<p>Can someone advise what’s the best way to have what I want implemented?  (Use DataLoader’s parallel loading, while simultaneously guaranteeing every batch loaded is entirely new.)</p>

username: SunQpark
cooked: <p>A solution worked for me was making a generator function using <code>itertools.repeat</code>.</p>
<pre><code class="lang-python">from itertools import repeat

def repeater(data_loader):
    for loader in repeat(data_loader):
        for data in loader:
            yield data
</code></pre>
<p>Then</p>
<pre><code class="lang-python">data_loader = DataLoader(dataset, ...)
data_loader = repeater(data_loader)

for data in data_loader:
    # train your model
</code></pre>

username: jteuwen
cooked: <p>As I ran into this as well, and likely many more will, this is my solution:<br>
If you look at <a href="https://pytorch.org/docs/stable/_modules/torch/utils/data/dataloader.html#_BaseDataLoaderIter" rel="nofollow noopener"><code>torch.utils.data.dataloader._BaseDataLoaderIter</code></a>` you will find:</p>
<pre><code class="lang-auto">def _next_index(self):
    return next(self._sampler_iter)  # may raise StopIteration
</code></pre>
<p>so a <code>StopIteration</code> is raised if your Sampler does so. You will also find in the <code>DataLoader</code> class itself:</p>
<pre><code class="lang-auto">    def __len__(self):
        if self._dataset_kind == _DatasetKind.Iterable:
            # (...) cut the comment here
            length = self._IterableDataset_len_called = len(self.dataset)
            return length
        else:
            return len(self._index_sampler)
</code></pre>
<p>that the length is defined by the <code>_index_sampler</code> if you do not have an <code>IterableDataset</code>. I think that the most elegant solution here is to have an infinite data sampler with no length (as the length is not required in the dataloader, see the comments in the code.</p>
<p>Something like this would be a good sampler:</p>
<pre><code class="lang-auto">class InfiniteSampler:
    def __init__(self, dataset_size):
        self.dataset_size = dataset_size

    def __iter__(self):
        yield from itertools.islice(self._infinite(), 0, None, 1) # Infinite iterator

    def _infinite(self):
        g = torch.Generator()
        while True:
            yield from torch.randperm(self.dataset_size, generator=g)
</code></pre>
<p>A similar approach is currently taken in <a href="https://github.com/facebookresearch/detectron2/blob/ef096f9b2fbedca335f7476b715426594673f463/detectron2/data/samplers/distributed_sampler.py" rel="nofollow noopener">Detectron2</a>.</p>

username: jonathanzjl
cooked: <p>Thanks! This is indeed the approach in detectron2. Just a kind reminder, there should probably be a seeding mechanism in the <strong>init</strong> and _infinite (something like the one in detectron2), otherwise the sampler will most likely generate the same pseudo-random sequence at each run. <img src="https://discuss.pytorch.org/images/emoji/apple/slightly_smiling_face.png?v=10" title=":slightly_smiling_face:" class="emoji" alt=":slightly_smiling_face:"></p>

username: ecobost
cooked: <p>islice is not needed. We can just return the generator (which is also an iterator); I think this version is easier to understand:</p>
<pre><code class="lang-auto">class InfiniteRandomSampler(data.Sampler):
    """ Return random indices from [0-n) infinitely.
    
    Arguments:
        dset_size (int): Size of the dataset to sample.
    """
    def __init__(self, dset_size):
        self.dset_size = dset_size

    def __iter__(self):
        # Create a random number generator (optional, makes the sampling independent of the base RNG)
        rng = torch.Generator()
        seed = torch.empty((), dtype=torch.int64).random_().item()
        rng.manual_seed(seed)
        
        return _infinite_generator(self.dset_size, rng)
    
    def __len__(self):
        return float('inf')
        
def _infinite_generator(n, rng):
    """ Inifinitely returns a number in [0, n)."""
    while True:
        yield from torch.randperm(n, generator=rng).tolist()
</code></pre>
<p>This will create an infinite dataloder but it won’t be “stateful”, i.e., if you stop sampling from it and then start sampling again, the sampling will start anew and forget which examples were already returned in the previous iterations. If you don’t want that, you’ll have to also use what <a class="mention" href="/u/sunqpark">@SunQpark</a> proposed (which also could drop the dependency on the itertools):</p>
<pre><code class="lang-auto">def make_infinite(dloader):
    while True:
        yield from dloader
</code></pre>