并行化 TensorFlow
为了扩展 TensorFlow 并行化的范围,我们还可以以分布式方式在完全不同的机器上从我们的图执行单独的操作。这个秘籍将告诉你如何。
做好准备
在 TensorFlow 发布几个月后,谷歌发布了分布式 TensorFlow,它是对 TensorFlow 生态系统的一次重大升级,并且允许在不同的工作机器上设置 TensorFlow 集群,并分享训练和评估的计算任务楷模。使用 Distributed TensorFlow 就像为工作器设置参数一样简单,然后为不同的工作器分配不同的工作。
在这个秘籍中,我们将建立两个本地工作器并将他们分配到不同的工作。
操作步骤
- 首先,我们加载 TensorFlow 并使用配置字典文件(ports
2222
和2223
)定义我们的两个本地 worker,如下所示:
import tensorflow as tf
# Cluster for 2 local workers (tasks 0 and 1):
cluster = tf.train.ClusterSpec({'local': ['localhost:2222', 'localhost:2223']})
- 现在,我们将两个 worker 连接到服务器并使用以下任务编号标记它们:
server = tf.train.Server(cluster, job_name="local", task_index=0)
server = tf.train.Server(cluster, job_name="local", task_index=1)
- 现在我们将让每个工作器完成一项任务。第一个工作器将初始化两个矩阵(每个矩阵将是 25 乘 25)。第二个工作器将找到所有元素的总和。然后,我们将自动分配两个总和的总和并打印输出,如下所示:
mat_dim = 25
matrix_list = {}
with tf.device('/job:local/task:0'):
for i in range(0, 2):
m_label = 'm_{}'.format(i)
matrix_list[m_label] = tf.random_normal([mat_dim, mat_dim])
# Have each worker calculate the sums
sum_outs = {}
with tf.device('/job:local/task:1'):
for i in range(0, 2):
A = matrix_list['m_{}'.format(i)]
sum_outs['m_{}'.format(i)] = tf.reduce_sum(A)
# Sum all the sums
summed_out = tf.add_n(list(sum_outs.values()))
with tf.Session(server.target) as sess:
result = sess.run(summed_out)
print('Summed Values:{}'.format(result))
- 输入上面的代码后,我们可以在命令提示符下运行以下命令:
$ python3 parallelizing_tensorflow.py
I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:197] Initialize GrpcChannelCache for job local -> {0 -> localhost:2222, 1 -> localhost:2223}
I tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc:206] Started server with target: grpc://localhost:2222
I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:197] Initialize GrpcChannelCache for job local -> {0 -> localhost:2222, 1 -> localhost:2223}
I tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc:206] Started server with target: grpc://localhost:2223
I tensorflow/core/distributed_runtime/master_session.cc:928] Start master session 252bb6f530553002 with config:
Summed Values:-21.12611198425293
工作原理
使用 Distributed TensorFlow 非常简单。您所要做的就是将工作者 IP 分配给具有名称的服务器。然后,可以手动或自动为操作员分配操作。