How to create a dataset and dataloader for whole images?

[malfonsoarquimea]

Hi! I am currently working on something where I train a model using individual pixels as data points. For evaluation, I need to generate results for a number of images, each composed of many pixels. For each pixel, I am computing some gradients and storing some things in memory, so a whole image has a memory footprint larger than my VRAM.
How would you suggest approaching this problem? A dataset/dataloader that loads individual pixels (like the ones I use for training) would be enough but I would need to reconstruct the images using all the pixels in multiple batches. Do you know of any solution that is more elegant?
I really appreciate any help you can provide.

[akihironitta]

Hi @malfonsoarquimea, sounds interesting!

so a whole image has a memory footprint larger than my VRAM.

I assume you get OOM during training. There're a number of techniques you could try to reduce the memory footprint:

• reduce the batch size and use gradient accumulation Trainer(accumulate_grad_batches=...)
• use amp Trainer(precision=16)
• use model parallelism (deepspeed, fairscale) when your model is very large: https://pytorch-lightning.readthedocs.io/en/stable/advanced/model_parallel.html

A dataset/dataloader that loads individual pixels (like the ones I use for training) would be enough but I would need to reconstruct the images using all the pixels in multiple batches. Do you know of any solution that is more elegant?

I would suggest loading only a subset of pixels in an image needed for each training_step. (Is this what you're doing already or trying to do next?)