RAM Memory overflow with GAN when using tensorflow.data

Ask Questions Forum: ask Machine Learning Questions to our readersCategory: TensorFlow/KerasRAM Memory overflow with GAN when using tensorflow.data
Vedant asked 1 month ago

Basically I’ve been training a GAN  and the program is crashing due to system memory overflow (not the GPU memory). as per my knowledge there is some method in tensorflow where we can directly load the data from disk (SSD) to the GPU by avoiding the RAM.  
the image size I’m using is 1024x1024x3 and using tensorflow.data  API to load the data.  
are you aware about how can I skip the data loading into RAM and directly load it into the GPU (If any such thing exists).  
I’ve tried Batch processing but unfortunately it didn’t work.    
you can find the code at https://github.com/JVedant/CycleGANs/blob/main/monet2photo.ipynb for reference.

1 Answers
Best Answer
Chris Staff answered 1 month ago

Okay, here is an idea.

1. Your training loop currently iterates over a Dataset that you construct using x_train, y_train, which you have loaded into memory. for image_x, image_y in tf.data.Dataset.zip((train_x, train_y)):
You report that this process fails, so:
2. What if you can construct a tf.data.Dataset object that is iterable but which does only load data into memory once it needs it?
That is, using a generator. Like this: https://www.machinecurve.com/index.php/2020/04/06/using-simple-generators-to-flow-data-from-file-with-keras/
This link however doesn’t work directly with your notebook because the link uses model.fit whereas you define a custom training loop with gradient tape. So:
3. Why don’t you create the dataset with a from_generator call, pointing towards a disk on your system where the data is located? https://www.tensorflow.org/api_docs/python/tf/data/Dataset#from_generator
The only thing you have to do is to manually download the cycle_gan/monet2photo data, unpack it, and construct the from_generator in such a way that it loads the data from file.

I think this can work.

from_generator takes any generator, even a simple generator –> here, my article can be useful, because it provides a standard def.