ОК ... так что с помощью моего коллеги у меня есть работоспособное решение. Ключ, по сути, является модификацией конфига. Конкретно как то так:
config.gpu_options.visible_device_list = '0'
Это гарантирует, что тензор потока видит только графический процессор 0.
На самом деле я смог запустить следующее:
#!/usr/bin/env python
import tensorflow as tf
with tf.device('/gpu:2'):
a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')
b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b')
c = tf.matmul(a, b)
# Creates a session with log_device_placement set to True.
config=tf.ConfigProto(log_device_placement=True)
config.gpu_options.visible_device_list = '0,1,3'
sess = tf.Session(config=config)
# Runs the op.
print(sess.run(c))
Обратите внимание, что этот код фактически указывает на работу на GPU 2 (который, как вы помните, является полным). Это важный момент ... графические процессоры перенумерованы в соответствии с visible_device_list, поэтому в приведенном выше коде, когда мы говорим «с gpu: 2», это относится к третьему графическому процессору в списке ('0,1,3 '), который на самом деле является графическим процессором 3. Это может укусить вас, если вы попробуете это:
#!/usr/bin/env python
import tensorflow as tf
with tf.device('/gpu:1'):
a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')
b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b')
c = tf.matmul(a, b)
# Creates a session with log_device_placement set to True.
config=tf.ConfigProto(log_device_placement=True)
config.gpu_options.visible_device_list = '1'
sess = tf.Session(config=config)
# Runs the op.
print(sess.run(c))
Проблема в том, что он ищет 2-й графический процессор в списке, но в видимом списке есть только один графический процессор. Вы получите следующую ошибку:
InvalidArgumentError (см. Выше для отслеживания): не удается назначить устройство для операции 'a': Операция была явно назначена для / device: GPU: 1, но доступны устройства [/ job: localhost / replica: 0 / task: 0 / устройство: ЦП: 0, / задание: локальный хост / реплика: 0 / задание: 0 / устройство: GPU: 0]. Убедитесь, что в спецификации устройства указано действительное устройство.
[[Узел: a = Constdtype = DT_FLOAT, значение = Тензор, _device = "/ устройство: GPU: 1"]]
Мне все еще кажется странным, что я должен указывать список графических процессоров, когда хочу запустить процессор. Я попытался использовать пустой список, и это не удалось, поэтому, если бы использовались все 4 графических процессора, я бы не нашел обходного пути. У кого-нибудь есть идея получше?