Функция config.set_visible_devices

Функция config.set_visible_devices применяется для ограничения набора физических устройств, доступных TensorFlow. Первым параметром передается список устройств или одно устройство, которое нужно сделать видимым. Вторым параметром можно передать тип устройства, например 'GPU' или 'CPU'. Если первым параметром передать пустой список, все устройства указанного типа станут невидимыми.

Функция особенно полезна при работе на машинах с несколькими GPU, когда нужно явно указать, какие из них будут использоваться. Настройку следует выполнять до создания тензоров и запуска вычислений.

Синтаксис

tf.config.set_visible_devices(devices, device_type=None)

Пример

Давайте посмотрим, какие физические устройства доступны в системе до изменения настроек:

import tensorflow as tf

devices = tf.config.list_physical_devices()
print(devices)

Результат выполнения кода:

[PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU'), PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]

Пример

Давайте сделаем видимым только CPU, скрыв все GPU от TensorFlow:

import tensorflow as tf

tf.config.set_visible_devices([], 'GPU')

devices = tf.config.get_visible_devices()
print(devices)

Результат выполнения кода:

[PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU')]

Пример

Давайте сделаем видимым конкретный GPU с индексом 0:

import tensorflow as tf

gpus = tf.config.list_physical_devices('GPU')
if gpus:
	tf.config.set_visible_devices(gpus[0], 'GPU')

devices = tf.config.get_visible_devices('GPU')
print(devices)

Результат выполнения кода:

[PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]

Пример

Давайте проверим, что вычисления выполняются на видимом устройстве:

import tensorflow as tf

tf.config.set_visible_devices([], 'GPU')

t = tf.constant([1, 2, 3, 4, 5])
print(t.device)

Результат выполнения кода:

"/job:localhost/replica:0/task:0/device:CPU:0"

Смотрите также

функцию config.list_physical_devices,
которая возвращает список физических устройств
функцию config.get_visible_devices,
которая возвращает список видимых устройств
функцию config.list_logical_devices,
которая возвращает список логических устройств
функцию experimental.set_memory_growth,
которая включает динамическое выделение памяти на GPU