Multi-process data loading with iterable dataset

username: mk6
cooked: <p>Hi,</p>
<p>I’m implementing the multi-process data loading logic for my own <code>Iterable</code> dataset. However I observed a strange behavior while playing with the <strong>second</strong> example in the doc <a href="https://pytorch.org/docs/stable/data.html#torch.utils.data.IterableDataset" rel="noopener nofollow ugc">here</a>, with implementing <code>worker_init_fn</code>.</p>
<p>The following is a code snippet to reproduce. <code>MyIterableDataset</code> and <code>worker_init_fn</code> are copied from the doc without any modification.</p>
<pre><code class="lang-auto">import torch
import math
from torch.utils.data import DataLoader


class MyIterableDataset(torch.utils.data.IterableDataset):
    def __init__(self, start, end):
        super(MyIterableDataset).__init__()
        assert end &gt; start, "this example code only works with end &gt;= start"
        self.start = start
        self.end = end

    def __iter__(self):
        return iter(range(self.start, self.end))


# Define a `worker_init_fn` that configures each dataset copy differently
def worker_init_fn(worker_id):
    worker_info = torch.utils.data.get_worker_info()
    dataset = worker_info.dataset  # the dataset copy in this worker process
    overall_start = dataset.start
    overall_end = dataset.end
    # configure the dataset to only process the split workload
    per_worker = int(math.ceil((overall_end - overall_start) / float(worker_info.num_workers)))
    worker_id = worker_info.id
    dataset.start = overall_start + worker_id * per_worker
    dataset.end = min(dataset.start + per_worker, overall_end)


if __name__ == '__main__':
    ds = MyIterableDataset(start=0, end=500)

    dl = DataLoader(
        dataset=ds, batch_size=100, num_workers=2, worker_init_fn=worker_init_fn,
    )

    for e in dl:
        print(e.shape)
</code></pre>
<p>Running this snippet gives the result:</p>
<pre><code class="lang-auto">torch.Size([100])
torch.Size([100])
torch.Size([100])
torch.Size([100])
torch.Size([50])
torch.Size([50])
</code></pre>
<p>We can see that the last <code>100</code> examples are split into two batches of <code>50</code> examples.</p>
<p>Similarly, when I changed <code>num_workers=3</code>, I got</p>
<pre><code class="lang-auto">torch.Size([100])
torch.Size([100])
torch.Size([100])
torch.Size([67])
torch.Size([67])
torch.Size([66])
</code></pre>
<p>Is this a bug or under expectation?</p>
<p>Thanks</p>

username: ptrblck
cooked: <p>Based on the code snippet the shapes seem to be expected.<br>
If you <code>print</code> the start and end indices from each worker you would see:</p>
<pre><code class="lang-python"># num_workers=2
worker 0, start 0, end 250
worker 1, start 250, end 500

# num_workers=3
worker 0, start 0, end 167
worker 1, start 167, end 334
worker 2, start 334, end 500
</code></pre>
<p>Each worker will thus yield full batches until the last one contains the remaining samples, as you’ve explicitly defined this logic in the <code>worker_init_fn</code> and <code>__iter__</code>.</p>

username: mk6
cooked: <p>Hi <a class="mention" href="/u/ptrblck">@ptrblck</a>,</p>
<p>Thanks for your help!</p>
<p>Each worker yields full batches first until none of them has a full batch, and then they yield respective incomplete batches even though virtually they can be combined into full batches, with <code>drop_last=False</code> in <code>Dataloader</code>.<br>
Is this what you mean?</p>
<p>My expectation was <code>Dataloader</code> would take the job of combining examples from different workers into complete batches when there are sufficient examples.</p>

username: ptrblck
cooked: <aside class="quote no-group" data-username="mk6" data-post="3" data-topic="151771">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/m/7993a0/48.png" class="avatar"> mk6:</div>
<blockquote>
<p>My expectation was <code>Dataloader</code> would take the job of combining examples from different workers into complete batches when there are sufficient examples.</p>
</blockquote>
</aside>
<p>I think this feature request is still not implemented and currently each worker creates the batch itself, so your explanation sounds right.</p>

username: mk6
cooked: <p>I see. Thanks for help!</p>