Vous pourriez utiliser tf.config.set_visible_devices
. Une fonction possible qui vous permet de définir si et quels GPU utiliser est:
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)
Supposons que vous soyez sur un système avec 4 GPU et que vous ne souhaitiez utiliser que deux GPU, celui avec id = 0
et celui avec id = 2
, alors la première commande de votre code, immédiatement après l'importation des bibliothèques, serait:
set_gpu([0, 2])
Dans votre cas, pour n'utiliser que le CPU, vous pouvez appeler la fonction avec une liste vide :
set_gpu([])
Par souci d'exhaustivité, si vous souhaitez éviter que l'initialisation du runtime n'alloue toute la mémoire du périphérique, vous pouvez utiliser tf.config.experimental.set_memory_growth
. Enfin, la fonction de gestion des appareils à utiliser, occupant dynamiquement la mémoire des GPU, devient:
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)