Memory Usage During running a Deep learning CNN Model in Colab

As explained in this post and my own observations, Tensorflow always tries to allocate the entire memory, no matter how small or big is your model. Unlike for example MXNet that only allocates enough memory to run the model.

Diving a little deeper, I learned that this is indeed the default
behaviour in Tensorflow: use all available RAM to speed things up.
Fair enough 🙂

You might think more memory allocation means faster training, but that’s not the case most of the times. You can restrict your TF memory usage, as shown in the following code:

import tensorflow as tf
from keras.backend.tensorflow_backend import set_session

config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.9
config.gpu_options.visible_device_list = "0"


Here is the Tensorflow documentation if you need more details of how to set restrictions on TF memory usage.

Leave a Comment