Google Cloud ML 엔진에 분산 된 Tensorflow 기기 배치

Question

Google 클라우드 ML 엔진에서 대규모 분산 Tensorflow 모델을 실행하고 있습니다. GPU가있는 머신을 사용하고 싶습니다. 내 그래프는 입력/데이터 판독기 기능과 계산 부분의 두 가지 주요 부분으로 구성됩니다.

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"으로 고정되어 있으며, 이는 기본적으로 "작업자"작업의 첫 번째 작업에 의해 관리되는 첫 번째 장치로 설정됩니다.

당신 임베디드 디바이스 블록을 사용하여 다른 장치 (또는 작업)로 고정 할 수 있습니다 :) 자신의 tf.device (지정

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