사용할 수 있습니다 tf.config.set_visible_devices
. 사용할 GPU를 설정할 수있는 기능 중 하나는 다음과 같습니다.
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)
4 개의 GPU가있는 시스템에 있고 두 개의 GPU, 즉를 사용하는 GPU와를 사용하는 GPU 만 사용하려는 경우 라이브러리를 가져온 직후 코드의 첫 번째 명령은 다음 id = 0
과 id = 2
같습니다.
set_gpu([0, 2])
귀하의 경우 CPU 만 사용하려면 빈 목록으로 함수를 호출 할 수 있습니다 .
set_gpu([])
완전성을 위해 런타임 초기화가 장치의 모든 메모리를 할당하지 않도록하려면 tf.config.experimental.set_memory_growth
. 마지막으로 GPU 메모리를 동적으로 점유하여 사용할 장치를 관리하는 기능은 다음과 같습니다.
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)