TL;DR: During on_train_epoch_start, I want to get the model's output on ALL the training data, as part of some pre-training calculations. I'm asking what the lightning-friendly way to do that is.
This is an odd question.
In my project, every 10 epochs I select a subset of the full training data, and train only on that subset. During part of the calculation of which subset to use, I compute the model's output on every datapoint in the train dataset.
My question is, what's the best way to do this in pytorch lightning? Currently I have a callback with an on_train_epoch_start hook. During this hook, the callback makes its own dataloader from trainer.datamodule.train_dataloader(), and manually iterates over the dataloader, computing the model outputs. That's not ideal, right?
This makes me run into problems with pytorch lightning. For instance, when training on the GPU, I get an error, since my callback is using its own dataloader, not the trainer's train_dataloader, and so it isn't on the GPU. However, I can't use the trainer's train_dataloader, since after my callback selects its subset, it changes the trainer's train_dataloader to be just that subset, instead of the full dataset.
I guess I have two main questions:
- Is there any way to avoid making a separate dataloader? Can I call a train loop somehow? Getting a models output on the full dataset seems like such a simple operation, I would think it'd be a one-liner.
- How can I get/use a dataloader that syncs with all of the Pytorch Lightning automatic modifications? (eg. GPU/CPU, dataloading workers, pin_memory)