Is there any way to disable prefetching of batches?

[jacobpbotha]

I noticed that one batch always gets prefetched, since I am using pytorch-lightning for Reinforcement Learning this is not desired behaviour.

Is there any way to have the trainer not do this?

[rohitgr7]

not an expert in Reinforcement learning, but is it something that can be configured inside native PyTorch Dataloader using prefetch_factor??
https://pytorch.org/docs/stable/data.html#torch.utils.data.DataLoader

[jacobpbotha]

I am pretty certain that it happens in the pytorch_lightning code. Just using the DataLoader normally gives me the expected behaviour (no prefetching), but once I pass it to the trainer it starts prefetching. I think it happens in the FitLoop.advance() code, where the dataloader gets wrapped with train_dataloader = self.trainer.data_connector.get_profiled_train_dataloader(train_dataloader), which wraps it with prefetch_iterator from trainer/supporters.py. I'm just wondering if there are any work arounds, becauses it means that I can't use the DataLoader interface in any of my RL code.

[rohitgr7]

cc @tchaton

[tchaton]

Dear @blacknfuzzy,

Yes, PyTorch Lightning is responsible to do the pre-fetching.
Could you explain us why you would need it to be prevent pre-fetching ?

PyTorch Lightning currently does the following in order to know if the batch is the latest and conclude properly an epoch.

Right now, the only work-around would be for you to replace the DataFetcher and keep track of latest.

Might be simpler if you could point us to some code, so we can brainstorm on the best way forward.

Best,
T.C

[jacobpbotha]

I need to prevent pre-fetching since I am generating the data during the training, and some algorithms require that the data is generated using the same model that is being updated. If there is pre-fetching that is no longer the case and any guarantee of correctness or convergence might no longer hold.

I can find the point in the code where it is wrapped in the pre-fetch iterator, but I don't really understand why or what it is used for at the moment. When I have the time I'll try to get a better understanding and see how I go at replacing the DataFetcher.

At the moment I'm simply using a dummy DataLoader (using range), and generating the data in the training_step function, but that means none of the callbacks or hooks actually have access to the batch.

This might also affect the Reinforcement Learning algorithms in lightning-bolts.

[tchaton]

This behaviour has always been present in Lightning, I am curious to hear about this only now :)
I would investigate deeper.

[mzweilin]

By defining __len__() to return a large-enough number, you can avoid pre-fetching in PL.

class Dataset:

  def __len__(self):
    return 65535
  
  def __iter__(self):
    return self
  
  def __next__(self):
    pass