Google Cloud MLエンジンにTensorflowデバイスを配置する

Question

Google Cloud MXエンジンで大規模な分散Tensorflowモデルを実行しています。私はGPUを搭載したマシンを使いたいです。 私のグラフは、入力/データリーダー機能と計算部分の2つの主要な部分で構成されています。

PSタスク、CPUの入力部分、GPUの計算部分に変数を配置したいと考えています。 関数tf.train.replica_device_setterは自動的に変数をPSサーバに配置します。

これは私のコードは次のようになります。

with tf.device(tf.train.replica_device_setter(cluster=cluster_spec)): 
    input_tensors = model.input_fn(...) 
    output_tensors = model.model_fn(input_tensors, ...) 

は、のようにreplica_device_setter()と一緒tf.device()を使用することが可能です：

with tf.device(tf.train.replica_device_setter(cluster=cluster_spec)): 
    with tf.device('/cpu:0') 
     input_tensors = model.input_fn(...) 
    with tf.device('/gpu:0') 
     tensor_dict = model.model_fn(input_tensors, ...) 

replica_divice_setter()を上書きし、変数が中に配置されていないことだろうPSサーバー？

さらに、クラスタ内のデバイス名はjob:master/replica:0/task:0/gpu:0のようなものなので、Tensorflow tf.device(whatever/gpu:0)にはどうすればいいですか？

Answer 1

tf.train.replica_device_setterブロック内の変数を超える操作は自動的に"/job:worker"に固定されます。これはデフォルトで「ワーカー」ジョブの最初のタスクによって管理される最初のデバイスになります。

あなたは組み込みデバイスのブロックを使って、別のデバイス（またはタスク）にそれらを固定することができます

with tf.device(tf.train.replica_device_setter(ps_tasks=2, ps_device="/job:ps", 
              worker_device="/job:worker")): 
    v1 = tf.Variable(1., name="v1") # pinned to /job:ps/task:0 (defaults to /cpu:0) 
    v2 = tf.Variable(2., name="v2") # pinned to /job:ps/task:1 (defaults to /cpu:0) 
    v3 = tf.Variable(3., name="v3") # pinned to /job:ps/task:0 (defaults to /cpu:0) 
    s = v1 + v2   # pinned to /job:worker (defaults to task:0/cpu:0) 
    with tf.device("/task:1"): 
    p1 = 2 * s   # pinned to /job:worker/task:1 (defaults to /cpu:0) 
    with tf.device("/cpu:0"): 
     p2 = 3 * s   # pinned to /job:worker/task:1/cpu:0