Tensorflow contrib.learn.Estimator multi-GPU

Question

マルチGPUトレーニングにcontrib.learn.Estimatorを使用するには、model_fnにGPU割り当てを指定しようとしています。擬似コードで

：

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

     return predictions, loss, train_op 

すべてがtf.device（ '/ GPU：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. 

私は、グラフに自分自身を悪用しているopを追加しているとは思っていませんが、Estimatorのスナップショット機能を通じて注入されています。

解決策はallow_soft_placement = Trueを設定して、GPU以外の機能がCPUに落ちるようにすることだと思いますが、contrib.learn.Estimatorを扱う際にどのように公開されているのかはわかりません。

通常、このオプションはConfigProto &に渡されますが、セッションの管理にはEstimatorの機能を使用しています。セッションの作成を制御する必要がありますか、これを達成するためにどこかのパラメータが欠けていますか？

事前にアドバイスをいただき、ありがとうございます。

Answer 1

EstimatorがTensorflow 1.0にcontribを残すのと同様に、これは修正されています。