Функция config.experimental.set_virtual_device_configuration

Функция config.experimental.set_virtual_device_configuration применяется к физическому устройству TensorFlow и позволяет создать на его основе виртуальные устройства с заданными ограничениями по памяти. Первым параметром функция принимает физическое устройство, вторым - список конфигураций виртуальных устройств. Каждая конфигурация создается через tf.config.experimental.VirtualDeviceConfiguration и может содержать параметр memory_limit.

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

Синтаксис

tf.config.experimental.set_virtual_device_configuration(
	device,
	virtual_devices
)

Пример

Давайте получим список физических устройств и выведем информацию о них:

import tensorflow as tf

physical_devices = tf.config.list_physical_devices('CPU')
print(physical_devices)

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

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

Пример

Давайте настроим два виртуальных устройства на базе одного физического CPU с ограничением памяти:

import tensorflow as tf

physical_devices = tf.config.list_physical_devices('CPU')

tf.config.experimental.set_virtual_device_configuration(
	physical_devices[0],
	[
		tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024),
		tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024)
	]
)

logical_devices = tf.config.list_logical_devices('CPU')
print(logical_devices)

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

[LogicalDevice(name='/device:CPU:0', device_type='CPU'), LogicalDevice(name='/device:CPU:1', device_type='CPU')]

Пример

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

import tensorflow as tf

physical_devices = tf.config.list_physical_devices('CPU')

tf.config.experimental.set_virtual_device_configuration(
	physical_devices[0],
	[
		tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024),
		tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024)
	]
)

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

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

tf.Tensor([11 12 13 14 15], shape=(5,), dtype=int32)

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

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