Not understanding how to make iterable dataset

EvanZ: <p>I’m having a hard time understanding the example iterable dataset in the docs, and how to apply it to my own code. The example code is:</p>
<pre><code class="lang-auto">&gt;&gt;&gt; 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):
...         worker_info = torch.utils.data.get_worker_info()
...         if worker_info is None:  # single-process data loading, return the full iterator
...             iter_start = self.start
...             iter_end = self.end
...         else:  # in a worker process
...             # split workload
...             per_worker = int(math.ceil((self.end - self.start) / float(worker_info.num_workers)))
...             worker_id = worker_info.id
...             iter_start = self.start + worker_id * per_worker
...             iter_end = min(iter_start + per_worker, self.end)
...         return iter(range(iter_start, iter_end))
...
</code></pre>
<p>In the <code>IterableDataset</code> I have created it loops through a set of files (the data is already randomly shuffled, so that is not an issue), and yields one row at a time:</p>
<pre><code class="lang-auto">    def __iter__(self) -&gt; Iterator[dict]:
        for file in self.files:
            if self.compressed:
                fopen = BZ2File(filename=file, mode='r')
            else:
                fopen = open(file=file, mode='r')

            with fopen as f:
                for row in f:
                    data = json.loads(row)
                    tokens = data['words']
                    tokens = tokens[:self.max_length]
                    random.shuffle(tokens)
                    indices, mask = self.tokens_to_indices(tokens)

                    item = {
                        'src': indices,
                        'mask': mask,
                        'label': data['label']
                    }

                    yield item
</code></pre>
<p>I’m not sure how it would make sense to return an iterator in this scenario. Any advice would be appreciated!</p>

actuallyaswin: <p>It’s fine to use yield inside the <code>__iter__</code> function. To instantiate the dataset and use it, you can do something similar to the following:</p>
<pre><code class="lang-auto">dataset = MyIterableDataset(...)
dataset = iter(dataset)  # makes it actually iterable

for n_iteration in range(100):
   batch_item = next(dataset)
   # batch_item is a dict with "src", "mask", and "label"
   ...
</code></pre>
<p>Note that your dataset will eventually “end” since you’re doing a finite for-loop over <code>self.files</code>. If you want your IterableDataset to continue indefinitely, just wrap everything with a while-loop at the top, e.g.:</p>
<pre><code class="lang-auto">    def __iter__(self) -&gt; Iterator[dict]:
        while True:
            for file in self.files:
                if self.compressed:
            ...
</code></pre>

EvanZ: <p>Hey <a class="mention" href="/u/actuallyaswin">@actuallyaswin</a> thanks for your response! So about the last point you made that eventually the dataset will “end”, I found actually that isn’t the case. In fact it keeps iterating through batches even after it should have already exhausted all the data. But…perhaps this is a behavior for <code>pytorch-lightning</code> which I’m using to train the model.</p>

EvanZ: <p>On a related note it looks like I can take advantage of multiple workers simply using islice when iterating over the files:</p>
<pre><code class="lang-auto">        worker_total_num = torch.utils.data.get_worker_info().num_workers
        worker_id = torch.utils.data.get_worker_info().id
        for file in self.files:
            if self.compressed:
                fopen = BZ2File(filename=file, mode='r')
            else:
                fopen = open(file=file, mode='r')

            with fopen as f:
                for row in islice(f,worker_id,None,worker_total_num):

</code></pre>