Limiting GPU memory use in Tensorflow

I am interested in deep learning, and even built my own PC recently with a GTX 1660 Super card so that I can do a bit of simple deep learning. But as I use the same PC for other stuff, I can't have the GPU using all the memory for deep learning. So here is what I do in Tensorflow to limit the amount of memory being used.

import tensorflow as tf

gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
  try:
    # Currently, memory growth needs to be the same across GPUs
    for gpu in gpus:
      tf.config.experimental.set_memory_growth(gpu, True)
    logical_gpus = tf.config.experimental.list_logical_devices('GPU')
    print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
  except RuntimeError as e:
    # Memory growth must be set before GPUs have been initialized
    print(e)

MEMORY_LIMIT=4096
gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
  # Restrict TensorFlow to only allocate 1GB of memory on the first GPU
  try:
    tf.config.experimental.set_virtual_device_configuration(
        gpus[0], [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=MEMORY_LIMIT)])
    logical_gpus = tf.config.experimental.list_logical_devices('GPU')
    print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPUs")
  except RuntimeError as e:
    # Virtual devices must be set before GPUs have been initialized
    print(e)

This sets memory use to grow as required, and limits the maximum memory use to whatever MEMORY_LIMIT (in megabytes) is set to.

 

Anyway, the 1660 Super is not exactly the best card for deep learning. The new RTX 30xx cards are definitely better and they are cheaper (relative to RTX 20xx cards). For those who can afford it, I think the RTX 30xx cards would be better for any serious work in deep learning. I am just using this budget build for my own learning.