Anda bisa menggunakan tf.config.set_visible_devices
. Satu kemungkinan fungsi yang memungkinkan Anda untuk mengatur jika dan GPU mana yang akan digunakan adalah:
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)
Misalkan Anda menggunakan sistem dengan 4 GPU dan Anda ingin menggunakan hanya dua GPU, satu dengan id = 0
dan satu dengan id = 2
, maka perintah pertama kode Anda, segera setelah mengimpor pustaka, adalah:
set_gpu([0, 2])
Dalam kasus Anda, untuk hanya menggunakan CPU, Anda dapat menjalankan fungsi dengan daftar kosong :
set_gpu([])
Untuk kelengkapannya, jika Anda ingin menghindari bahwa inisialisasi waktu proses akan mengalokasikan semua memori pada perangkat, Anda dapat menggunakan tf.config.experimental.set_memory_growth
. Terakhir, fungsi untuk mengelola perangkat mana yang akan digunakan, menempati memori GPU secara dinamis, menjadi:
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)