Podrías usar tf.config.set_visible_devices
. Una función posible que le permite establecer si y qué GPU usar es:
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)
Suponga que está en un sistema con 4 GPU y desea usar solo dos GPU, la que tiene id = 0
y la que tiene id = 2
, entonces el primer comando de su código, inmediatamente después de importar las bibliotecas, sería:
set_gpu([0, 2])
En su caso, para usar solo la CPU, puede invocar la función con una lista vacía :
set_gpu([])
Para completar, si desea evitar que la inicialización en tiempo de ejecución asigne toda la memoria en el dispositivo, puede usar tf.config.experimental.set_memory_growth
. Finalmente, la función para administrar qué dispositivos usar, ocupando la memoria de la GPU de forma dinámica, se convierte en:
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)