Tensorflow contrib.learn.Estimator multi-GPU

Question

Чтобы использовать contrib.learn.Estimator для обучения с несколькими GPU, я пытаюсь указать назначения графического процессора в моей модели_fn.

В псевдокоде:

def model_fn(X, y): 
    with tf.device('/gpu:1'): 
     ... various tensorflow ops for model ... 

     return predictions, loss, train_op 

Все отлично работает без tf.device ('/ ГПУ: 1') вызов, но с этим я сталкиваюсь следующее сообщение об ошибке:

InvalidArgumentError (see above for traceback): Cannot assign a device to 
node 'save/ShardedFilename_1': Could not satisfy explicit device 
specification '/device:GPU:1' because no supported kernel 
for GPU devices is available. 

Я не считаю, что я добавляю оскорбительный оператор к графику, но скорее, что он вводится через функциональность моментального снимка Estimator.

Я считаю, что решение должно установить allow_soft_placement = True, чтобы функции не GPU попадали в CPU, но для меня не очевидно, как это проявляется при работе с contrib.learn.Estimator.

Я вижу, что этот параметр обычно задается в ConfigProto &, переданном сеансу, но я использовал функциональность Estimator для управления сеансом для меня. Должен ли я взять контроль над созданием сеанса, или я пропустил какой-то параметр для этого?

Большое спасибо за любой совет.

Answer 1

Наряду с оценщиком, оставляющим вклад в Tensorflow 1.0, это исправлено.