Custom dataloader behaviour

[juneskiafc]

Hi,
I'm looking to change the dataloader every epoch.
The context is multiple instance learning, where top k patches are selected from each large image, and these patches are different across epochs.
To do this, first I need to do one complete pass over all patches in an image. I do this with a separate dataloader, in the on_epoch_start hooks.
Then, after I do this inference pass, using its results I want to reload the dataloader inside the LightningModule so that regular training_step, validation_step, test_step will read the top k patches only. However, I don't know of an easy way to do this with pytorch-lightning.
Does anyone have any ideas? It would be greatly appreciated.

[Borda]

@juneskiafc can you implement the dataloader update/replacement in the on_epoch_end hook or is there some limitation?

[juneskiafc]

I've currently solved this problem by just passing in separate dataloaders to the model (in addition to the datamodule). The problem in trying to reassign the dataloader in the model hook was that, although I reassigned the dataloader, train_step didn't seem to use the updated dataloader.