How to use DataLoader with IterableDataset

username: kfshr
cooked: <p>I would like to use <code>IterableDataset</code> to create an infinite dataset that I can pass to <code>DataLoader</code>.</p>
<p>I tried two approaches and would like to know which one should be preferred or if there is a better solution for an infinite stream of data in Pytorch.</p>
<p>I also have the problem that changing the <code>batch_size</code> argument has no effect. What do I do wrong here?</p>
<p>What about other parameters such as <code>pin_memory</code>, <code>shuffle</code> etc? Can I simply ignore them when working with an infinite dataset?</p>
<p>Here is a minimal working example:</p>
<pre><code class="lang-auto">import torch
from torch.utils.data.dataloader import DataLoader
from torch.utils.data import IterableDataset


class DataStream1(IterableDataset):

    def __init__(self) -&gt; None:
        super().__init__()
        self.size_input = 4
        self.size_output = 2

    def generate(self):
        x = torch.rand(self.size_input) 
        y = torch.rand(self.size_output) 
        yield x, y 

    def __iter__(self):
        return iter(self.generate())


class Stream:

    def __init__(self):
        self.size_input = 4
        self.size_output = 2

    def __next__(self):
        x = torch.rand(self.size_input) 
        y = torch.rand(self.size_output) 
        return x, y 


class DataStream2(IterableDataset):

    def __init__(self) -&gt; None:
        super().__init__()

    def __iter__(self):
        return Stream() 


def test(data_set):
    train_loader = DataLoader(
        dataset=data_set,
        pin_memory=False,
        shuffle=False,
        batch_size=2,
        num_workers=1,
    )
    for epoch in range(2):
        print(f"\n{epoch = }")
        for x, y in train_loader:
            print(f"{x = }")
            print(f"{y = }")


def main():

    # Ignores batch size, but stops after two epochs.
    data_stream_1 = DataStream1()
    test(data_stream_1)

    # Ignores batch size, does exit the train_loader loop.
    data_stream_2 = DataStream2()
    test(data_stream_2)


if __name__ == "__main__":
    main()
</code></pre>

username: arunppsg
cooked: <p>For an infinite streaming dataset, how about something like this:</p>
<pre><code class="lang-auto">class DataStream1(IterableDataset):

    def __init__(self) -&gt; None:
        super().__init__()
        self.size_input = 4
        self.size_output = 2

    def generate(self):
        while True:
            x = torch.rand(self.size_input)
            y = torch.rand(self.size_output)
            yield x, y

    def __iter__(self):
        return iter(self.generate())

dataset = DataStream1()

train_loader = DataLoader(dataset=dataset)

for i, data in enumerate(train_loader):
    print (i, data)
</code></pre>

username: kfshr
cooked: <p>Thank you for your response.</p>
<p>I don’t see how your answer is different from my <code>DataStream1</code> class. I also would like to know how I can define a batch size? Maybe you can say something more about that using your approach. Thank you.</p>

username: arunppsg
cooked: <p>In the <code>generate</code> function, I have added a <code>while True</code> condition which can yield an infinite stream of data. To have a batch size, it can be updated as:</p>
<pre><code class="lang-auto">    def generate(self):
        while True:
            x = torch.rand(self.batch_size, self.size_input)
            y = torch.rand(self.batch_size, self.size_output)
            yield x, y
</code></pre>