[python] tensorflow에서 현재 사용 가능한 GPU를 얻는 방법은 무엇입니까?

분산 TensorFlow를 사용할 계획이 있으며 TensorFlow가 훈련 및 테스트에 GPU를 사용할 수 있다는 것을 알았습니다. 클러스터 환경에서 각 머신은 0 또는 1 이상의 GPU를 가질 수 있으며 가능한 많은 머신에서 GPU로 TensorFlow 그래프를 실행하고 싶습니다.

tf.Session()TensorFlow를 실행할 때 아래와 같이 로그 메시지에서 GPU에 대한 정보를 제공 한다는 것을 알았습니다 .

I tensorflow/core/common_runtime/gpu/gpu_init.cc:126] DMA: 0
I tensorflow/core/common_runtime/gpu/gpu_init.cc:136] 0:   Y
I tensorflow/core/common_runtime/gpu/gpu_device.cc:838] Creating TensorFlow device (/gpu:0) -> (device: 0, name: GeForce GTX 1080, pci bus id: 0000:01:00.0)

내 질문은 TensorFlow에서 현재 사용 가능한 GPU에 대한 정보를 어떻게 얻습니까? 로그에서 GPU 정보를로드 할 수는 있지만보다 정교하고 프로그래밍 방식으로 수행하고 싶습니다. CUDA_VISIBLE_DEVICES 환경 변수를 사용하여 의도적으로 GPU를 제한 할 수도 있으므로 OS 커널에서 GPU 정보를 얻는 방법을 알고 싶지 않습니다.

요컨대, 기계에 두 개의 GPU가 있으면 tf.get_available_gpus()반환 되는 함수 ['/gpu:0', '/gpu:1']가 필요합니다. 이것을 어떻게 구현할 수 있습니까?

답변

from tensorflow.python.client import device_lib

def get_available_gpus():
    local_device_protos = device_lib.list_local_devices()
    return [x.name for x in local_device_protos if x.device_type == 'GPU']

from tensorflow.python.client import device_lib

device_lib.list_local_devices()

tf.test.is_gpu_available()

tf.test.gpu_device_name()

import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
    print("Name:", gpu.name, "  Type:", gpu.device_type)

Name: /physical_device:GPU:0   Type: GPU
Name: /physical_device:GPU:1   Type: GPU

gpus = tf.config.list_physical_devices('GPU')

import subprocess

n = str(subprocess.check_output(["nvidia-smi", "-L"])).count('UUID')

import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
    print("Name:", gpu.name, "  Type:", gpu.device_type)

    gpus = tf.config.list_physical_devices('GPU')