[Python] Clearing tf.data.Dataset from GPU memory

I'm running into an issue when implementing a training loop that uses a tf.data.Dataset as input to a Keras model. My dataset has an element spec of the following format:

({'data': TensorSpec(shape=(15000, 1), dtype=tf.float32), 'index': TensorSpec(shape=(2,), dtype=tf.int64)}, TensorSpec(shape=(1,), dtype=tf.int32))


So, basically, each sample is structured as tuple (x, y), in which x has the structure of a dict containing two tensors, one of data with shape (15000, 1), and the other an index of shape (2,) (the index is not used during training), and y is a single label.

The tf.data.Dataset is created using dataset = tf.data.Dataset.from_tensor_slices((X, y)), where X is a dict of two keys:

• data: an np array of shape (200k, 1500, 1), index with
• index: an np array of shape (200k, 2)

and y is a single array of shape (200k, 1)

My dataset has about 200k training samples (after running undersampling) and 200k validation samples.

Right after calling tf.data.Dataset.from_tensor_slices I noticed a spike in GPU memory usage, with about 16GB being occupied after creating the training tf.Dataset, and 16GB more after creating the validation tf.Dataset.

After creating of the tf.Dataset, I run a few operations (e.g. shuffle, batching, and prefetching), and call model.fit. My model has about 500k trainable parameters.

The issue I'm running into is after fitting the model. I need to run inference on some additional data, so I create a new tf.Dataset with this data, again using tf.Dataset.from_tensor_slices. However, I noticed the training and validation tf.Dataset still reside in GPU memory, which causes my script to break with an out of memory problem for the new tf.Dataset I want to run inference on.

I tried calling del on the two tf.Dataset, and subsequently calling gc.collect(), but I believe that will only clear RAM, not GPU memory. Also, I tried disabling some operations I apply, such as prefetch, and also playing with the batch size, but none of that worked. I also tried calling keras.backend.clear_session(), but it also did not work to clear GPU memory. I also tried importing cuda from numba, but due to my install I cannot use it to clear memory. Is there any way for me to clear the tf.data.Dataset from GPU memory?

Continue reading...