Vị trí thiết bị phân phối Tensorflow trong công cụ Google Cloud ML

Question

Tôi đang chạy mô hình Tensorflow phân tán lớn trong công cụ google cloud ML. Tôi muốn sử dụng máy có GPU. Biểu đồ của tôi bao gồm hai phần chính là chức năng đầu vào/đầu đọc dữ liệu và phần tính toán.

Tôi muốn đặt các biến trong tác vụ PS, phần đầu vào trong CPU và phần tính toán trên GPU. Chức năng tf.train.replica_device_setter sẽ tự động đặt các biến trong máy chủ PS.

Đây là những gì mã của tôi trông giống như:

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

Có thể sử dụng tf.device() cùng với replica_device_setter() như trong:

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, ...) 

Liệu replica_divice_setter() được ghi đè và biến không được đặt trong máy chủ PS?

Hơn nữa, vì tên thiết bị trong cụm là một cái gì đó như job:master/replica:0/task:0/gpu:0 làm cách nào để nói với Tensorflow tf.device(whatever/gpu:0)?

Answer 1

Mọi hoạt động, ngoài biến số, trong khối tf.train.replica_device_setter được tự động ghim vào "/job:worker", sẽ được mặc định thành thiết bị đầu tiên được quản lý bởi tác vụ đầu tiên trong công việc "công nhân".

Bạn thể ghim chúng vào một thiết bị khác (hoặc công việc) bằng cách sử dụng khối thiết bị nhúng:

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