本節中的代碼大量使用『TensorFlow』分布式訓練_其一_邏輯梳理中介紹的概念，是成熟的多機分布式訓練樣例
一、基本概念
Cluster、Job、task概念：三者可以簡單的看成是層次關系，task可以看成每台機器上的一個進程，多個task組成job；job又有：ps、worker兩種，分別用於參數服務、計算服務，組成cluster。
同步更新
各個用於並行計算的電腦，計算完各自的batch 后，求取梯度值，把梯度值統一送到ps服務機器中，由ps服務機器求取梯度平均值，更新ps服務器上的參數。
如下圖所示，可以看成有四台電腦，第一台電腦用於存儲參數、共享參數、共享計算，可以簡單的理解成內存、計算共享專用的區域，也就是ps job；另外三台電腦用於並行計算的，也就是worker task。
這種計算方法存在的缺陷是：每一輪的梯度更新，都要等到A、B、C三台電腦都計算完畢后，才能更新參數，也就是迭代更新速度取決與A、B、C三台中，最慢的那一台電腦，所以采用同步更新的方法，建議A、B、C三台的計算能力都不想。
異步更新
ps服務器收到只要收到一台機器的梯度值，就直接進行參數更新，無需等待其它機器。這種迭代方法比較不穩定，收斂曲線震動比較厲害，因為當A機器計算完更新了ps中的參數，可能B機器還是在用上一次迭代的舊版參數值。
二、抽象接口
1、定義分布式集群對象
tf.train.ClusterSpec
# coding=utf-8 # 多台機器，每台機器有一個顯卡、或者多個顯卡，這種訓練叫做分布式訓練 import tensorflow as tf # 現在假設我們有A、B、C、D四台機器，首先需要在各台機器上寫一份代碼，並跑起來，各機器上的代碼內容大部分相同 # 除了開始定義的時候，需要各自指定該台機器的task之外。
# 以機器A為例子，A機器上的代碼如下： cluster=tf.train.ClusterSpec({ "worker": [ "A_IP:2222", # 格式 IP地址：端口號，第一台機器A的IP地址 ,在代碼中需要用這台機器計算的時候，就要定義：/job:worker/task:0 "B_IP:1234" # 第二台機器的IP地址 /job:worker/task:1 "C_IP:2222" # 第三台機器的IP地址 /job:worker/task:2 ], "ps": [ "D_IP:2222", # 第四台機器的IP地址 對應到代碼塊：/job:ps/task:0 ]})
然后我們需要寫四分代碼，這四分代碼文件大部分相同，但是有幾行代碼是各不相同的(可以通過命令行參數使得字面相同各自選擇if分支不同)。
2、在各台機器上，定義server
比如A機器上的代碼server要定義如下：
server=tf.train.Server(cluster,job_name='worker',task_index=0) # 找到‘worker’名字下的，task0，也就是機器A
3、在代碼中，指定device
device不僅可以指定卡，還可以將計算圖中的不同節點指定不同機器
with tf.device('/job:ps/task:0'): # 參數定義在機器D上
w=tf.get_variable('w',(2,2),tf.float32,initializer=tf.constant_initializer(2))
b=tf.get_variable('b',(2,2),tf.float32,initializer=tf.constant_initializer(5))
with tf.device('/job:worker/task:0/cpu:0'): # 在機器A cpu上運行
addwb=w+b
with tf.device('/job:worker/task:1/cpu:0'): # 在機器B cpu上運行
mutwb=w*b
with tf.device('/job:worker/task:2/cpu:0'): # 在機器C cpu上運行
divwb=w/b
不過在深度學習訓練圖計算中，對於每個worker task來說，計算任務都是相同的，所以我們會把所有圖計算、變量定義等代碼，都寫到下面這個語句下：
with tf.device(tf.train.replica_device_setter(worker_device='/job:worker/task:indexi',cluster=cluster)):
函數replica_deviec_setter會自動把變量參數定義部分定義到ps服務中(如果ps有多個任務，那么自動分配)。
下面舉個例子，假設現在有兩台機器A、B，A用於計算服務，B用於參數服務，那么代碼如下：
# 上面是因為worker計算內容各不相同，不過再深度學習中，一般每個worker的計算內容是一樣的，
# 因為都是計算神經網絡的每個batch的前向傳導，所以一般代碼是重用的
import tensorflow as tf
# 現在假設我們有A、B台機器，首先需要在各台機器上寫一份代碼，並跑起來，各機器上的代碼內容大部分相同
# ，除了開始定義的時候，需要各自指定該台機器的task之外。以機器A為例子，A機器上的代碼如下：
cluster=tf.train.ClusterSpec({
"worker": [
"192.168.11.105:1234", # 格式 IP地址：端口號，在代碼中需要用這台機器計算的時候，就要定義：/job:worker/task:0
],
"ps": [
"192.168.11.130:2223" # 第四台機器的IP地址 對應到代碼塊：/job:ps/task:0
]})
# 不同的機器，下面這一行代碼各不相同，server可以根據job_name、task_index兩個參數，查找到集群cluster中對應的機器
isps=False
if isps:
server=tf.train.Server(cluster,job_name='ps',task_index=0) # 找到‘ps’名字下的，task0，也就是機器A
server.join()
else:
server=tf.train.Server(cluster,job_name='worker',task_index=0) # 找到‘worker’名字下的，task0，也就是機器B
with tf.device(tf.train.replica_device_setter(worker_device='/job:worker/task:0',cluster=cluster)):
w=tf.get_variable('w',(2,2),tf.float32,initializer=tf.constant_initializer(2))
b=tf.get_variable('b',(2,2),tf.float32,initializer=tf.constant_initializer(5))
addwb=w+b
mutwb=w*b
divwb=w/b
saver = tf.train.Saver()
summary_op = tf.merge_all_summaries()
init_op = tf.initialize_all_variables()
sv = tf.train.Supervisor(init_op=init_op, summary_op=summary_op, saver=saver)
with sv.managed_session(server.target) as sess:
while 1:
print sess.run([addwb,mutwb,divwb])
把該代碼在機器A上運行，程序會進入等候狀態，等候用於ps參數服務的機器啟動，才會運行。
因此接着我們需要在機器B上運行如下代碼：
# 上面是因為worker計算內容各不相同，不過再深度學習中，一般每個worker的計算內容是一樣的，
# 因為都是計算神經網絡的每個batch前向傳導，所以一般代碼是重用的
#coding=utf-8
#多台機器，每台機器有一個顯卡、或者多個顯卡，這種訓練叫做分布式訓練
import tensorflow as tf
# 現在假設我們有A、B、C、D四台機器，首先需要在各台機器上寫一份代碼，並跑起來，各機器上的代碼內容大部分相同
# ，除了開始定義的時候，需要各自指定該台機器的task之外。以機器A為例子，A機器上的代碼如下：
cluster=tf.train.ClusterSpec({
"worker": [
"192.168.11.105:1234", # 格式 IP地址：端口號，在代碼中需要用這台機器計算的時候，就要定義：/job:worker/task:0
],
"ps": [
"192.168.11.130:2223" # 第四台機器的IP地址 對應到代碼塊：/job:ps/task:0
]})
# 不同的機器，下面這一行代碼各不相同，server可以根據job_name、task_index兩個參數，查找到集群cluster中對應的機器
isps=True
if isps:
server=tf.train.Server(cluster,job_name='ps',task_index=0) # 找到‘ps’名字下的，task0，也就是機器A
server.join()
else:
server=tf.train.Server(cluster,job_name='worker',task_index=0) # 找到‘worker’名字下的，task0，也就是機器B
with tf.device(tf.train.replica_device_setter(worker_device='/job:worker/task:0',cluster=cluster)):
w=tf.get_variable('w',(2,2),tf.float32,initializer=tf.constant_initializer(2))
b=tf.get_variable('b',(2,2),tf.float32,initializer=tf.constant_initializer(5))
addwb=w+b
mutwb=w*b
divwb=w/b
saver = tf.train.Saver()
summary_op = tf.merge_all_summaries()
init_op = tf.initialize_all_variables()
sv = tf.train.Supervisor(init_op=init_op, summary_op=summary_op, saver=saver)
with sv.managed_session(server.target) as sess:
while 1:
print sess.run([addwb,mutwb,divwb])
附錄1：分布式訓練需要熟悉的函數
tf.train.Servertf.train.Supervisortf.train.SessionManagertf.train.ClusterSpectf.train.replica_device_settertf.train.MonitoredTrainingSessiontf.train.MonitoredSessiontf.train.SingularMonitoredSessiontf.train.Scaffoldtf.train.SessionCreatortf.train.ChiefSessionCreatortf.train.WorkerSessionCreator
附錄2：分布式訓練例程
tensorflow/tools/dist_test/python/mnist_replica.py
# Copyright 2016 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Distributed MNIST training and validation, with model replicas.
A simple softmax model with one hidden layer is defined. The parameters
(weights and biases) are located on one parameter server (ps), while the ops
are executed on two worker nodes by default. The TF sessions also run on the
worker node.
Multiple invocations of this script can be done in parallel, with different
values for --task_index. There should be exactly one invocation with
--task_index, which will create a master session that carries out variable
initialization. The other, non-master, sessions will wait for the master
session to finish the initialization before proceeding to the training stage.
The coordination between the multiple worker invocations occurs due to
the definition of the parameters on the same ps devices. The parameter updates
from one worker is visible to all other workers. As such, the workers can
perform forward computation and gradient calculation in parallel, which
should lead to increased training speed for the simple model.
"""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import math
import sys
import tempfile
import time
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data
flags = tf.app.flags
flags.DEFINE_string("data_dir", "/tmp/mnist-data",
"Directory for storing mnist data")
flags.DEFINE_boolean("download_only", False,
"Only perform downloading of data; Do not proceed to "
"session preparation, model definition or training")
flags.DEFINE_integer("task_index", None,
"Worker task index, should be >= 0. task_index=0 is "
"the master worker task the performs the variable "
"initialization ")
flags.DEFINE_integer("num_gpus", 1, "Total number of gpus for each machine."
"If you don't use GPU, please set it to '0'")
flags.DEFINE_integer("replicas_to_aggregate", None,
"Number of replicas to aggregate before parameter update "
"is applied (For sync_replicas mode only; default: "
"num_workers)")
flags.DEFINE_integer("hidden_units", 100,
"Number of units in the hidden layer of the NN")
flags.DEFINE_integer("train_steps", 200,
"Number of (global) training steps to perform")
flags.DEFINE_integer("batch_size", 100, "Training batch size")
flags.DEFINE_float("learning_rate", 0.01, "Learning rate")
flags.DEFINE_boolean(
"sync_replicas", False,
"Use the sync_replicas (synchronized replicas) mode, "
"wherein the parameter updates from workers are aggregated "
"before applied to avoid stale gradients")
flags.DEFINE_boolean(
"existing_servers", False, "Whether servers already exists. If True, "
"will use the worker hosts via their GRPC URLs (one client process "
"per worker host). Otherwise, will create an in-process TensorFlow "
"server.")
flags.DEFINE_string("ps_hosts", "localhost:2222",
"Comma-separated list of hostname:port pairs")
flags.DEFINE_string("worker_hosts", "localhost:2223,localhost:2224",
"Comma-separated list of hostname:port pairs")
flags.DEFINE_string("job_name", None, "job name: worker or ps")
FLAGS = flags.FLAGS
IMAGE_PIXELS = 28
def main(unused_argv):
mnist = input_data.read_data_sets(FLAGS.data_dir, one_hot=True)
if FLAGS.download_only:
sys.exit(0)
if FLAGS.job_name is None or FLAGS.job_name == "":
raise ValueError("Must specify an explicit `job_name`")
if FLAGS.task_index is None or FLAGS.task_index == "":
raise ValueError("Must specify an explicit `task_index`")
print("job name = %s" % FLAGS.job_name)
print("task index = %d" % FLAGS.task_index)
# 解析集群參數
# ps作業的ip端口，可以有多個ps作業，之間用逗號分割
ps_spec = FLAGS.ps_hosts.split(",")
# worker作業的ip端口，可以有多個worker作業，之間用逗號分割
worker_spec = FLAGS.worker_hosts.split(",")
# Get the number of workers.
num_workers = len(worker_spec)
# 分布式集群對象，字典形式接收作業類型對應的任務主機&端口
cluster = tf.train.ClusterSpec({"ps": ps_spec, "worker": worker_spec})
if not FLAGS.existing_servers:
# Not using existing servers. Create an in-process server.
# 任務內部服務器（ps或者worker的上層抽象）的抽象對象，接收集群對象，接收當前任務信息
server = tf.train.Server(
cluster, job_name=FLAGS.job_name, task_index=FLAGS.task_index)
# 如果本次程序運行於ps任務，則啟動監聽（join持續等待，不會返回）
if FLAGS.job_name == "ps":
server.join()
# 選擇本任務所處GPU
if FLAGS.num_gpus > 0:
# Avoid gpu allocation conflict: now allocate task_num -> #gpu
# for each worker in the corresponding machine
gpu = (FLAGS.task_index % FLAGS.num_gpus)
worker_device = "/job:worker/task:%d/gpu:%d" % (FLAGS.task_index, gpu)
elif FLAGS.num_gpus == 0:
# Just allocate the CPU to worker server
cpu = 0
worker_device = "/job:worker/task:%d/cpu:%d" % (FLAGS.task_index, cpu)
""" 構建網絡 """
# The device setter will automatically place Variables ops on separate
# parameter servers (ps). The non-Variable ops will be placed on the workers.
# The ps use CPU and workers use corresponding GPU
# 設備設定器將自動將變量OPS放置在單獨的參數服務器（PS）上。不可變的OPS將放在worker身上。
with tf.device(
# 設備放置器，可以返回被tf.device接受的設備名稱
tf.train.replica_device_setter(
worker_device=worker_device,
ps_device="/job:ps/cpu:0",
cluster=cluster)):
global_step = tf.Variable(0, name="global_step", trainable=False)
# Variables of the hidden layer
hid_w = tf.Variable(
tf.truncated_normal(
[IMAGE_PIXELS * IMAGE_PIXELS, FLAGS.hidden_units],
stddev=1.0 / IMAGE_PIXELS),
name="hid_w")
hid_b = tf.Variable(tf.zeros([FLAGS.hidden_units]), name="hid_b")
# Variables of the softmax layer
sm_w = tf.Variable(
tf.truncated_normal(
[FLAGS.hidden_units, 10],
stddev=1.0 / math.sqrt(FLAGS.hidden_units)),
name="sm_w")
sm_b = tf.Variable(tf.zeros([10]), name="sm_b")
# Ops: located on the worker specified with FLAGS.task_index
x = tf.placeholder(tf.float32, [None, IMAGE_PIXELS * IMAGE_PIXELS])
y_ = tf.placeholder(tf.float32, [None, 10])
hid_lin = tf.nn.xw_plus_b(x, hid_w, hid_b)
hid = tf.nn.relu(hid_lin)
y = tf.nn.softmax(tf.nn.xw_plus_b(hid, sm_w, sm_b))
cross_entropy = -tf.reduce_sum(y_ * tf.log(tf.clip_by_value(y, 1e-10, 1.0)))
opt = tf.train.AdamOptimizer(FLAGS.learning_rate)
# 如果使用同步訓練機制
if FLAGS.sync_replicas:
# 如果並行副本數(期望)沒有指定
if FLAGS.replicas_to_aggregate is None:
# 勒令並行數等於worker數
replicas_to_aggregate = num_workers
else:
# 用戶指定了就用指定的
replicas_to_aggregate = FLAGS.replicas_to_aggregate
# 同步優化器，接收本地優化器
opt = tf.train.SyncReplicasOptimizer(
opt,
replicas_to_aggregate=replicas_to_aggregate,
total_num_replicas=num_workers,
name="mnist_sync_replicas")
# 優化器或者同步優化器單步優化節點（注意此句在if外）
train_step = opt.minimize(cross_entropy, global_step=global_step)
# 如果是worker，則編號0的worker設置為chief worker
is_chief = (FLAGS.task_index == 0)
# 同步訓練機制下的初始化操作
if FLAGS.sync_replicas:
# local_step初始化（chief_worker會改寫此句，所以實際上本句針對非chief_worker）
local_init_op = opt.local_step_init_op
if is_chief:
# chief_worker使用的時global_step，也需要初始化
local_init_op = opt.chief_init_op
# 為未初始化的Variable初始化
ready_for_local_init_op = opt.ready_for_local_init_op
# Initial token and chief queue runners required by the sync_replicas mode
# 同步標記隊列實例
chief_queue_runner = opt.get_chief_queue_runner()
# 同步標記隊列初始值設定
sync_init_op = opt.get_init_tokens_op()
# 全局變量初始化
init_op = tf.global_variables_initializer()
train_dir = tempfile.mkdtemp()
if FLAGS.sync_replicas:
# 管理同步訓練相關操作
sv = tf.train.Supervisor(
is_chief=is_chief,
logdir=train_dir,
init_op=init_op,
local_init_op=local_init_op,
ready_for_local_init_op=ready_for_local_init_op,
recovery_wait_secs=1,
global_step=global_step)
else:
# 管理異步訓練相關操作
sv = tf.train.Supervisor(
is_chief=is_chief,
logdir=train_dir,
init_op=init_op,
recovery_wait_secs=1,
global_step=global_step)
# 配置分布式會話
# 沒有可用GPU時使用CPU
# 不打印設備放置信息
# 過濾未綁定在ps或者worker的操作
sess_config = tf.ConfigProto(
allow_soft_placement=True,
log_device_placement=False,
device_filters=["/job:ps",
"/job:worker/task:%d" % FLAGS.task_index])
# chief會初始化所有worker的會話，否則等待chief返回會話
# The chief worker (task_index==0) session will prepare the session,
# while the remaining workers will wait for the preparation to complete.
if is_chief:
print("Worker %d: Initializing session..." % FLAGS.task_index)
else:
print("Worker %d: Waiting for session to be initialized..." %
FLAGS.task_index)
if FLAGS.existing_servers:
server_grpc_url = "grpc://" + worker_spec[FLAGS.task_index]
print("Using existing server at: %s" % server_grpc_url)
sess = sv.prepare_or_wait_for_session(server_grpc_url, config=sess_config)
else:
sess = sv.prepare_or_wait_for_session(server.target, config=sess_config)
print("Worker %d: Session initialization complete." % FLAGS.task_index)
# 同步更新模式的chief worker
if FLAGS.sync_replicas and is_chief:
# Chief worker will start the chief queue runner and call the init op.
# 初始化同步標記隊列
sess.run(sync_init_op)
# 啟動相關線程，運行各自服務
sv.start_queue_runners(sess, [chief_queue_runner])
# Perform training
time_begin = time.time()
print("Training begins @ %f" % time_begin)
local_step = 0
while True:
# Training feed
batch_xs, batch_ys = mnist.train.next_batch(FLAGS.batch_size)
train_feed = {x: batch_xs, y_: batch_ys}
_, step = sess.run([train_step, global_step], feed_dict=train_feed)
local_step += 1
now = time.time()
print("%f: Worker %d: training step %d done (global step: %d)" %
(now, FLAGS.task_index, local_step, step))
if step >= FLAGS.train_steps:
break
time_end = time.time()
print("Training ends @ %f" % time_end)
training_time = time_end - time_begin
print("Training elapsed time: %f s" % training_time)
# Validation feed
val_feed = {x: mnist.validation.images, y_: mnist.validation.labels}
val_xent = sess.run(cross_entropy, feed_dict=val_feed)
print("After %d training step(s), validation cross entropy = %g" %
(FLAGS.train_steps, val_xent))
if __name__ == "__main__":
tf.app.run()