You could use tf.config.set_visible_devices
. One possible function that allows you to set if and which GPUs to use is:
import tensorflow as tf
def set_gpu(gpu_ids_list):
gpus = tf.config.list_physical_devices('GPU')
if gpus:
try:
gpus_used = [gpus[i] for i in gpu_ids_list]
tf.config.set_visible_devices(gpus_used, 'GPU')
logical_gpus = tf.config.experimental.list_logical_devices('GPU')
print(len(gpus), "Physical GPUs,", len(logical_gpus), "Logical GPU")
except RuntimeError as e:
# Visible devices must be set before GPUs have been initialized
print(e)
Suppose you are on a system with 4 GPUs and you want to use only two GPUs, the one with id = 0
and the one with id = 2
, then the first command of your code, immediately after importing the libraries, would be:
set_gpu([0, 2])
In your case, to use only the CPU, you can invoke the function with an empty list:
set_gpu([])
For completeness, if you want to avoid that the runtime initialization will allocate all memory on the device, you can use tf.config.experimental.set_memory_growth
.
Finally, the function to manage which devices to use, occupying the GPUs memory dynamically, becomes:
import tensorflow as tf
def set_gpu(gpu_ids_list):
gpus = tf.config.list_physical_devices('GPU')
if gpus:
try:
gpus_used = [gpus[i] for i in gpu_ids_list]
tf.config.set_visible_devices(gpus_used, 'GPU')
for gpu in gpus_used:
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 GPU")
except RuntimeError as e:
# Visible devices must be set before GPUs have been initialized
print(e)