Dataloader starving the gpu
Hey,
Thank you for this amazing library !
I'm using pytorch_lightning to train a segmentation model on 3D images. Augmentation on these images is quite slow, mostly because I do full volume elastic transforms, which takes ~2s per image on a single cpu.
I'm running a large unet with precision 16 and amp_level 'O2', I'm using pytorch 1.4, pytorch_lightning 7.1.
I noticed that these augmentations do lead to an underusage of the GPU which seems to be waiting for the data: not surprising if my augmentation is slow. But in fact, the time it takes to run an epoch is (much) larger than the time to run an epoch without augmentation + the time to do the augmentation.
To measure that, I changed my dataset so that it 'preloads' all the augmentations before the epoch starts i.e. loads in ram memory the tensors ready to be fed to the network, and then just return these tensors in the get_item method of my dataset. In this setting, I'm roughly 50% faster than originally, when doing the augmentation in the get_item method !
Is this expected behaviour ? Are the dataloaders pre-computing the get_item so that the GPU is fed as fast as possible ? More generally, what do you advise, when augmentation is slow, to best feed the gpu ?
Additional question: I couldn't make pytorch_lightning profiler work, I tried profiler=True in the trainer declaration or to set it to the BaseProfiler or to AdvancedProfiler and nothing seems to happen. Is this the right way to use it ?
Best,
Maxime