NotImplementedError after torch.utils.data.random_split

username: pytorchnewb
cooked: <p>I have an iterable dataset (and dataloader for it) which works fine. But after splitting it using <code>torch.utils.data.random_split</code>, and making the split datasets into a dataloader I’m getting a <code>NotImplementedError</code>.</p>
<p>Here’s a small example to reproduce the error:</p>
<pre><code class="lang-auto">class TestDataset(torch.utils.data.IterableDataset):
    def __init__(self):
        super().__init__()
        self.n = 10
        
    def __len__(self):
        return self.n
    
    def __iter__(self):
        for i in range(self.n):
            yield i
            
test_dataset = TestDataset()
test_dataset1, test_dataset2 = torch.utils.data.random_split(test_dataset, lengths=[8, 2])
test_dataloader1 = torch.utils.data.DataLoader(test_dataset1, batch_size=2)
next(iter(test_dataloader1))  # NotImplementedError
</code></pre>
<p>Why am I getting this error and how can I split an Iterable style dataset then?</p>

username: ptrblck
cooked: <p>The <code>Subset</code> tries to index the dataset, which doesn’t work for <code>IterableDataset</code>s:</p>
<pre><code class="lang-python">test_dataset = TestDataset()
test_dataset[0]
&gt; NotImplementedError
</code></pre>
<p>Since you already seem to define a<code>__len__</code> method I would suggest to use the map-style <code>torch.utils.data.Dataset</code> and define the <code>__getitem__</code> method, which would work using <code>random_split</code>.</p>

username: pytorchnewb
cooked: <p>Unfortunately that’s not possible. I can calculate the <code>len</code> but can’t keep all the data in memory to make it indexable. Is there no way to split this kind of a dataset? For example in tensorflow you can <code>.take</code> and <code>.skip</code> to generate a train and validation set.</p>

username: ptrblck
cooked: <p>You don’t have to preload the data and can lazily load it as is described in <a href="https://pytorch.org/tutorials/beginner/data_loading_tutorial.html">this tutorial</a>.</p>

username: pytorchnewb
cooked: <p>I don’t have a traditional dataset so unfortunately I can’t do that. It has to be generated on the fly and can’t be indexed into.</p>
<p>I’m just gonna hack something manually since I really need to get this done fast. Was hoping there would be some kind of a <code>take/skip</code> API since that should work regardless of iterable data or map-style.</p>

username: ptrblck
cooked: <p>In that case I don’t really understand the overall use case.<br>
<code>IterableDataset</code> is used for data “streams”, i.e. <code>Dataset</code>s, which cannot directly be indexed.<br>
The map-style datasets define a <code>__getitem___</code> as well as a length, so that indexing via <code>dataset[index]</code> would load the corresponding sample. <code>random_split</code> uses this property to split the indices into the desired ranges and returns the <code>Subset</code>s.<br>
However, based on your description you cannot calculate the length and don’t have a <code>__getitem__</code> (so the <code>IterableDataset</code> seems to be the right choice). Nevertheless you would like to “split” this dataset using indices, which doesn’t fit your described use case, since you can’t use indices to load samples.</p>