How to set dynamic memory growth on TF 2.1?

Question

With previous versions of tensorflow+keras I was able to set an 'allow_growth' option and view realtime memory usage with nvidia-smi. Otherwise it will all gett allocated immediately by the process. Now, using tf.keras in tensorflow 2.1 I can't find a way to do this. Any help is appreciated!

Answer 1

In case you have several GPUs, you will allow memory growth only for the first GPU.

physical_devices = tf.config.list_physical_devices('GPU') 
tf.config.experimental.set_memory_growth(physical_devices[0], True)

If you want to do it for all GPUs you need to set it for every instance.

physical_devices = tf.config.list_physical_devices('GPU') 
for gpu_instance in physical_devices: 
    tf.config.experimental.set_memory_growth(gpu_instance, True)