Функция device

Функция device создает контекстный менеджер, который указывает TensorFlow, на каком устройстве выполнять операции. Первым параметром функция принимает имя устройства в виде строки, например '/CPU:0' или '/GPU:0'. Внутри блока with все создаваемые тензоры и операции размещаются на выбранном устройстве.

Синтаксис

tf.device(device_name)

Пример

Давайте создадим тензор на устройстве CPU:

import tensorflow as tf

with tf.device('/CPU:0'):
	t = tf.constant([1, 2, 3, 4, 5])

print(t)

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

tf.Tensor([1 2 3 4 5], shape=(5,), dtype=int32)

Пример

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

import tensorflow as tf

with tf.device('/CPU:0'):
	t = tf.constant([1, 2, 3, 4, 5])

print(t.device)

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

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

Пример

Давайте выполним сложение тензоров внутри контекста устройства:

import tensorflow as tf

a = tf.constant([[1, 2, 3], [4, 5, 6]])
b = tf.constant([[1, 1, 1], [1, 1, 1]])

with tf.device('/CPU:0'):
	res = a + b

print(res)

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

tf.Tensor(
[[2 3 4]
 [5 6 7]], shape=(2, 3), dtype=int32)

Пример

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

import tensorflow as tf

with tf.device('/CPU:0'):
	v = tf.Variable([1, 2, 3, 4, 5])

print(v.numpy())

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

[1 2 3 4 5]

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

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