title: Implementing an “infinite loop” Dataset & DataLoader in PyTorch
body: <p>I’d like to implement an infinite loop Dataset &amp; DataLoader. Here’s what I tried:</p>
  
  <pre class="lang-py prettyprint-override"><code>class Infinite(Dataset):
      def __len__(self):
          return HPARAMS.batch_size
  #         return 1&lt;&lt;30 # This causes huge memory usage.
      def __getitem__(self, idx):
          """Randomly generates one new example."""
          return sample_func_to_be_parallelized()
  
  infinite_loader = DataLoader(
      dataset=Infinite(), 
      batch_size=HPARAMS.batch_size, 
      num_workers=16,
      worker_init_fn=lambda worker_id: np.random.seed(worker_id),  
  )
  
  while True:
      for idx, data in enumerate(infinite_loader):
          # forward + backward on "data"
  </code></pre>
  
  <p>As you can see, the main challenge here is the <code>__len()__</code> method. If I put a large enough number there, like 1&lt;&lt;30, the symptom is memory usage will JUMP TO 10+GB on the first iteration of train loop. After a while the workers are killed presumably due to OOM.</p>
  
  <p>If I put a small number there, like 1 or BATCH_SIZE, the sampled “data” in the train loop will be periodically duplicated. This is not what I want as I’d like new data to be generated &amp; trained on at every iteration.</p>
  
  <p>I’m guessing the culprit of the excessive memory usage is somewhere in the stack, a bunch of things are cached. Upon a casual look at the Python side of things I can’t pinpoint where.</p>
  
  <p>Can someone advise what’s the best way to have what I want implemented? (Use DataLoader’s parallel loading, while simultaneously guaranteeing every batch loaded is entirely new.)</p>