Enumerate(dataloader) slow

vibe2: <p>Calls to enumerate(dataloader) are quite slow in my project. I’m loading images stored in LMDB format, and I have multiple LMDBs that I call ConcatDataset on to create the final dataset. I noticed that reducing the num_workers from say 6 to 1 or 2 reduces the proportion of time spent on enumerate(), though that slows the loading of images. I was previously passing in a WeightedRandomSampler() into the dataloader and shuffle=False, but when I tried shuffle=True and sampler=None, that also reduced the amount of slowness around enumerate(). Despite these changes, about 20-30% of the training time goes to waiting for enumerate(dataloader) to start each epoch. My code looks like this:</p>
<pre><code class="lang-auto">for epoch in range(start_epoch, total_epochs+1):
    for _,train_data in enumerate(train_loader):
        #do work
</code></pre>
<p>I have confirmed that the hanging/slowness is due to enumerate and not the work being done during each epoch. Looking at the documentation (<a href="https://pytorch.org/docs/stable/data.html" class="inline-onebox" rel="noopener nofollow ugc">torch.utils.data — PyTorch 2.1 documentation</a>) I can see that the slowness is likely due to a large amount of work being done in each call of enumerate():</p>
<blockquote>
<p>In this mode, each time an iterator of a <a href="https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader" rel="noopener nofollow ugc"> <code>DataLoader</code> </a> is created (e.g., when you call  <code>enumerate(dataloader)</code> ),  <code>num_workers</code>  worker processes are created. At this point, the  <code>dataset</code> ,  <code>collate_fn</code> , and  <code>worker_init_fn</code>  are passed to each worker, where they are used to initialize, and fetch data. This means that dataset access together with its internal IO, transforms (including  <code>collate_fn</code> ) runs in the worker process. … Workers are shut down once the end of the iteration is reached, or when the iterator becomes garbage collected.</p>
</blockquote>
<p>My question is, what can we do to eliminate the slowness around enumerate? Is it possible to keep those processes generated in the first call to enumerate and pass them a different shuffling of the data? Or any other ideas?</p>
<p>Other info: I’m using the Apex mixed precision package, training on V100 GPUs, and using torch.distributed.DataParallel() on 4 GPUs. Trying apex.parallel.DistributedDataParallel did not significantly reduce the slowness. The hanging is still significant when training on one or two GPUs.</p>

ptrblck: <p>If your <code>Dataset.__init__</code> method is slow due to some heavy data loading, you would see the slowdown in each new creation of the workers.<br>
The recreation of the workers might yield a small slowdown, but should be negligible, if you are using lazy loading and don’t need a lot of resources in the <code>__init__</code> method.</p>
<p>Could you check, which operations are used in the <code>__init__</code>?</p>

vibe2: <p>Here’s what my <strong>init</strong> looks like for the dataset class:</p>
<pre><code class="lang-auto">   def __init__(self, opt):
        self.GT_env = lmdb.open(opt['dataroot_GT'], readonly=True, lock=False, readahead=False, meminit=True)
        keys = pickle.load(open(os.path.join(opt['dataroot_GT'], '_keys_cache.p'), "rb"))
        self.paths_GT = sorted([key for key in keys if not key.endswith('.meta')])

        self.LQ_env = lmdb.open(opt['dataroot_LQ'], readonly=True, lock=False, readahead=False, meminit=True)
        keys = pickle.load(open(os.path.join(opt['dataroot_LQ'], '_keys_cache.p'), "rb"))
        self.paths_LQ = sorted([key for key in keys if not key.endswith('.meta')])
</code></pre>
<p>It takes .005 seconds to run <strong>init</strong> per dataset, and I concatenate 8 datasets to make the final dataset. Let me know if you have any suggestions. Thanks for looking!</p>

rwightman: <p><a class="mention" href="/u/vibe2">@vibe2</a> A researcher submitted a solution to this issue to my codebase a while back, apparently it has also been discussed in a pytorch issue/PR for sometime</p>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/rwightman/pytorch-image-models/blob/d72ac0db259275233877be8c1d4872163954dfbb/timm/data/loader.py#L209-L238" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/rwightman/pytorch-image-models/blob/d72ac0db259275233877be8c1d4872163954dfbb/timm/data/loader.py#L209-L238" target="_blank" rel="nofollow noopener">rwightman/pytorch-image-models/blob/d72ac0db259275233877be8c1d4872163954dfbb/timm/data/loader.py#L209-L238</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="209" style="counter-reset: li-counter 208 ;">
<li>class MultiEpochsDataLoader(torch.utils.data.DataLoader):</li>
<li>
</li><li>    def __init__(self, *args, **kwargs):</li>
<li>        super().__init__(*args, **kwargs)</li>
<li>        self._DataLoader__initialized = False</li>
<li>        self.batch_sampler = _RepeatSampler(self.batch_sampler)</li>
<li>        self._DataLoader__initialized = True</li>
<li>        self.iterator = super().__iter__()</li>
<li>
</li><li>    def __len__(self):</li>
<li>        return len(self.batch_sampler.sampler)</li>
<li>
</li><li>    def __iter__(self):</li>
<li>        for i in range(len(self)):</li>
<li>            yield next(self.iterator)</li>
<li>
</li><li>
</li><li>class _RepeatSampler(object):</li>
<li>    """ Sampler that repeats forever.</li>
<li>
</li></ol></code></pre>

  This file has been truncated. <a href="https://github.com/rwightman/pytorch-image-models/blob/d72ac0db259275233877be8c1d4872163954dfbb/timm/data/loader.py#L209-L238" target="_blank" rel="nofollow noopener">show original</a>

  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>And the PR that was submitted with some info about the solution from the author: <a href="https://github.com/rwightman/pytorch-image-models/pull/140" rel="nofollow noopener">https://github.com/rwightman/pytorch-image-models/pull/140</a></p>

vibe2: <p>Works great. Thank you!</p>

toliveira: <p>One alternative is to use the parameter <code>persistent_workers</code> available in <a href="https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader" rel="noopener nofollow ugc">DataLoader</a><br>
<strong>persistent_workers</strong> (<a href="https://docs.python.org/3/library/functions.html#bool" rel="noopener nofollow ugc"><em>bool</em></a><em>,</em> <em>optional</em>) – If <code>True</code>, the data loader will not shutdown the worker processes after a dataset has been consumed once. This allows to maintain the workers Dataset instances alive. (default: <code>False</code>)</p>

bryant1410: <p>Note it isn’t the same. This implementation maintains the workers alive but also pre-fetches for the next epoch as the sampler is infinite. My understanding is that PyTorch’s persistent-workers feature won’t pre-fetch for the next epoch as the sampler would be empty.</p>