Możesz użyć tf.config.set_visible_devices
. Jedną z możliwych funkcji, która pozwala ustawić, czy i które GPU mają być używane, to:
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)
Załóżmy, że pracujesz w systemie z 4 procesorami graficznymi i chcesz używać tylko dwóch procesorów graficznych, jednego z id = 0
i tego z id = 2
, a pierwsze polecenie kodu, natychmiast po zaimportowaniu bibliotek, brzmiałoby:
set_gpu([0, 2])
W twoim przypadku, aby używać tylko procesora, możesz wywołać funkcję z pustą listą :
set_gpu([])
Aby uzyskać kompletność, jeśli chcesz uniknąć sytuacji, w której inicjalizacja środowiska wykonawczego przydzieli całą pamięć na urządzeniu, możesz użyć tf.config.experimental.set_memory_growth
. Wreszcie, funkcja zarządzania urządzeniami do wykorzystania, dynamicznie zajmującymi pamięć GPU, staje się:
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)