Как получить текущие доступные графические процессоры в тензорном потоке?

У меня есть план использования распределенного TensorFlow, и я увидел, что TensorFlow может использовать графические процессоры для обучения и тестирования. В среде кластера каждая машина может иметь 0 или 1 или более 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 из журнала, но я хочу сделать это более сложным, программным способом. Я также мог бы намеренно ограничить использование графических процессоров с помощью переменной среды CUDA_VISIBLE_DEVICES, поэтому я не хочу знать способ получения информации о GPU из ядра ОС.

Короче говоря, мне нужна функция типа tf.get_available_gpus(), которая вернет ['/gpu:0', '/gpu:1'], если в машине есть два графических процессора. Как это реализовать?

Ответ 1

Существует недокументированный метод device_lib.list_local_devices(), который позволяет вам отображать устройства, доступные в локальном процессе. ( NB Как недокументированный метод, это зависит от несовместимых изменений.) Функция возвращает список DeviceAttributes buffer buffer объектов. Вы можете извлечь список имен строковых устройств для устройств GPU следующим образом:

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']

Обратите внимание, что (по крайней мере, до TensorFlow 1.4) при вызове device_lib.list_local_devices() будет запущен некоторый код инициализации, который по умолчанию будет распределять всю память GPU на всех устройствах (проблема GitHub). Чтобы этого избежать, сначала создайте сеанс с явно небольшим per_process_gpu_fraction или allow_growth=True, чтобы предотвратить выделение всей выделенной памяти. Подробнее см. этот вопрос.

Ответ 2

Вы можете проверить список всех устройств, используя следующий код:

from tensorflow.python.client import device_lib

device_lib.list_local_devices()

Ответ 3

В утилите test также есть метод. Так что все, что нужно сделать, это:

tf.test.is_gpu_available()

и/или

tf.test.gpu_device_name()

Посмотрите документы Tensorflow для аргументов.

Ответ 4

Принятый ответ дает вам количество графических процессоров, но также выделяет всю память на этих графических процессорах. Вы можете избежать этого, создав сеанс с фиксированной нижней памятью перед вызовом device_lib.list_local_devices(), что может быть нежелательным для некоторых приложений.

В итоге я использовал nvidia-smi, чтобы получить количество графических процессоров без выделения памяти на них.

import subprocess

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

Ответ 5

Помимо превосходного объяснения, сделанного Мрри, где он предложил использовать device_lib.list_local_devices() я могу показать вам, как вы можете проверять информацию, связанную с GPU, из командной строки.

Поскольку в настоящее время только Nvidia GPUS работают для платформ NN, ответ охватывает только их. У Nvidia есть страница, где они документируют, как вы можете использовать интерфейс файловой системы /proc для получения во время выполнения информации о драйвере, любых установленных графических картах NVIDIA и статусе AGP.

/proc/driver/nvidia/gpus/0..N/information

Предоставьте информацию о каждом из установленных графических адаптеров NVIDIA (название модели, IRQ, версия BIOS, тип шины). Обратите внимание, что версия BIOS доступна только во время работы X.

Таким образом, вы можете запустить это из командной строки cat/proc/driver/nvidia/gpus/0/information и посмотреть информацию о вашем первом GPU. Это легко запустить из Python, а также вы можете проверить второй, третий, четвертый GPU, пока он не выйдет из строя.

Определенно, ответ Mrry более надежный, и я не уверен, будет ли мой ответ работать на машинах, отличных от Linux, но эта страница Nvidia предоставляет другую интересную информацию, о которой мало кто знает.

Ответ 6

В TensorFlow 2.0 вы можете использовать tf.config.experimental.list_physical_devices('GPU'):

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

Вы можете узнать больше на этой странице документации.

Ответ 7

Следующее работает в tenorflow 2.0:

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