search

clementchadebec / benchmark_VAE

Unifying Variational Autoencoder (VAE) implementations in Pytorch (NeurIPS 2022)
Apache License 2.0
1.82k stars 163 forks source link

pin_memory in DataLoader #153

Open psteinb opened 2 months ago

psteinb commented 2 months ago

Thanks for the library. I'm currently exploring it for somewhat larger images (greyscale, 3400x3400).

Is your feature request related to a problem? Please describe. I discovered that the BaseTrainer does create a DataLoader, does not touch the pin_memory flag. For larger data like mine, it would be great if this parameter can be turned on OR the use gets an option to turn it on when configuring the Trainer.

Describe the solution you'd like Either add a flag to BaseTrainerConfig or add something like dataloader_kwargs to it for use specific settings.

Describe alternatives you've considered Providing my data in a custom DataLoader. That appears to work, but I am not sure what side effects this can have (thinking of DDP parallelisation scheme).

Additional context Much of the benchmarks are performed with "small" data. RGB images with up to 64x64 pixels. And I also have the feeling that this was a predominant assumption when coding up the base encoders and such (lots of Linear layers).

psteinb commented 2 months ago

Btw, I am happy to send a PR on this, but I'd feel more comfortable doing so after we discussed the matter.