When does a Pytorch Dataset worker pull new data?

Sam_Lerman: <p>Does a <code>torch.utils.data.Dataset</code> worker initiate fetching when its iterator is called? <code>__iter__</code>? <code>__next__</code>? Once the DataLoader collates the batch, does the worker automatically start pulling the next one or does it wait idly until the iterator is called again?</p>
<p>TLDR: At what point does it fetch? What triggers the fetch?</p>

ptrblck: <p>The workers should start loading the data once the iterator is created and preload the <code>prefetch_factor*num_workers</code> batches. Afterwards each <code>next</code> call will consume a batch and let the corresponding worker load the next batch.<br>
Here is a small example showing this behavior:</p>
<pre><code class="lang-python">class MyDataset(Dataset):
    def __init__(self):
        self.data = torch.randn(512, 1)
        
    def __getitem__(self, idx):
        worker_info = torch.utils.data.get_worker_info()
        print("loading idx {} from worker {}".format(
            idx, worker_info.id))
        x = self.data[idx]
        return x
    
    def __len__(self):
        return len(self.data)


dataset = MyDataset()
loader = DataLoader(dataset, num_workers=8, batch_size=2)
loader_iter = iter(loader) # preloading starts here
# with the default prefetch_factor of 2, 2*num_workers=16 batches will be preloaded
# the max index printed by __getitem__ is thus 31 (16*batch_size=32 samples loaded)

data = next(loader_iter) # this will consume a batch and preload the next one from a single worker to fill the queue
# batch_size=2 new samples should be loaded
</code></pre>

Sam_Lerman: <p>One more question while you’re here…</p>
<p>Suppose one of the workers modifies their own data instance:</p>
<pre><code class="lang-auto">class MyDataset(Dataset):
    def __init__(self):
        self.data = torch.randn(512, 1)
        
    def __getitem__(self, idx):
        worker_info = torch.utils.data.get_worker_info()
        if worker_info.id == 2:
            self.data[7, :] = 5  
        x = self.data[idx]
        return x
    
    def __len__(self):
        return len(self.data)
</code></pre>
<p>Now from outside the worker, is there a way for me to get that worker’s version of the data? As in:</p>
<pre><code class="lang-auto">loader = DataLoader(dataset, num_workers=8, batch_size=2)
data_from_worker = loader.get_worker(id=2).data
</code></pre>
<p>Is this possible?</p>

ptrblck: <p>I’m not aware of a mechanism to access the workers directly besides their information in the <code>Dataset</code>.</p>

Vinayaka_Hegde: <p>Thanks for the answer above. I wanted to know where exactly in the dataloader.py “the <code>next</code> call would consume the current batch and load the next batch” ?</p>
<p>As far as I understood, the <code>__next__</code> function calls the `_next_data.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/1/e1968a362d763310a85fab97b6d1e4cf3e305ab8.png" data-download-href="https://discuss.pytorch.org/uploads/default/e1968a362d763310a85fab97b6d1e4cf3e305ab8" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/e/1/e1968a362d763310a85fab97b6d1e4cf3e305ab8.png" alt="image" data-base62-sha1="wbDXAnlGpygSZ9tozcpiIE3aiPS" width="690" height="174" data-dominant-color="F7F6F8"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">721×182 6.88 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>The <code>_next-data</code> calls the <code>_process_data</code><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/f/0f62dbf8ce02d421d79a686455bfac56540b4843.png" data-download-href="https://discuss.pytorch.org/uploads/default/0f62dbf8ce02d421d79a686455bfac56540b4843" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/0/f/0f62dbf8ce02d421d79a686455bfac56540b4843.png" alt="image" data-base62-sha1="2c6XXmqEf9aDOSOGMiVRbbqewjp" width="646" height="500" data-dominant-color="FAFAF9"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1082×837 40.6 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>and finally, in the <code>process_data</code> method, the <code>try_put_index</code> is being called which fetches the next index from sampler and gets the next active worker and assigns the batch to the worker.<br>
<img src="https://discuss.pytorch.org/uploads/default/original/3X/4/1/414e14085532909c4be282bf825fae63d2939e93.png" alt="image" data-base62-sha1="9jIjGkiMFSki0GMnr4tasJaXBED" width="602" height="176"></p>
<p>On the other hand, for consuming the current index, the worker.py has a line to perform a <code>indexqueue.get</code> that gets the current index from the queue and loads the corresponding batch using <code>fetcher.fetch</code>.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/c/6c27968163a4daa5ac5c5e894f80bc9e7101dfea.png" data-download-href="https://discuss.pytorch.org/uploads/default/6c27968163a4daa5ac5c5e894f80bc9e7101dfea" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/6/c/6c27968163a4daa5ac5c5e894f80bc9e7101dfea.png" alt="image" data-base62-sha1="fqMnMfaRk7Wdo6gqbwg4BhimV1U" width="690" height="43" data-dominant-color="FAF8F8"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">894×56 2.4 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Can you please confirm if I understand this properly ?</p>

ptrblck: <p>The description of the call chain sounds reasonable based on the posted screenshots of the code, but to be sure I would recommend to verify it by either stepping through it with a debugger or by e.g. adding debug <code>print</code> statements to see how the execution is done.<br>
Also related to <a href="https://discuss.pytorch.org/t/guidance-to-understand-the-dataloader-architecture-in-a-minute-level/170563">this thread</a>.</p>

aoot: <p>Does this mean that <code>loader_iter = iter(loader)</code> only needs to be called once in a training loop because each time a batch is consumed via the call <code>next(loader_iter)</code>, a batch will be pre-loaded to maintain the number of pre-loaded batches?</p>

ptrblck: <p>Yes, that’s right but note that you would need to recreate the <code>iterator</code> once it’s exhausted unless you are setting <code>persistent_workers=True</code>.</p>

Laubeee: <p>I read this and thought: when the iterator is exhausted it will automatically start again from index 0 when <code>persistent_workers=True</code> ? The docs do not say anything about that…</p>
<p>i am currently calling <code>iter(dataloader)</code> at the start of every epoch, and I was wondering if in the background the next prefetching happens or not, but it sounds like I don’t need to worry about re-creating the iterator before processing the last batch of one epoch? Do I not need to re-create it at all?</p>