深入浅出TensroFlow（七）TensorFlow计算加速

在前面的文章中介绍了使用TensorFlow实现各种深度学习的算法。然而要将深度学习应用到实际问题中，一个非常大的问题在于训练深度学习模型需要的计算量太大。比如要将Inception-v3模型在单机单卡上训练到78%的正确率需要将近半年的时间，这样的训练速度是完全无法应用到实际生产中的。为了加速训练过程，本文将介绍如何通过TensorFlow利用GPU或/和分布式计算进行模型训练。

TensorFlow使用GPU

TensorFlow程序可以通过tf.device函数来指定运行每一个操作的设备，这个设备可以是本地的CPU或者GPU，也可以是某一台远程的服务器。TensorFlow会给每一个可用的设备一个名称，tf.device函数可以通过设备的名称来指定执行运算的设备。比如CPU在TensorFlow中的名称为/cpu:0。

在默认情况下，即使机器有多个CPU，TensorFlow也不会区分它们，所有的CPU都使用/cpu:0作为名称。而一台机器上不同GPU的名称是不同的，第n个GPU在TensorFlow中的名称为/gpu:n。比如第一个GPU的名称为/gpu:0，第二个GPU名称为/gpu:1，以此类推。

TensorFlow提供了一个快捷的方式来查看运行每一个运算的设备。在生成会话时，可以通过设置log_device_placement参数来打印运行每一个运算的设备。下面的程序展示了如何使用log_device_placement这个参数。

import tensorflow as tf

a = tf.constant([1.0, 2.0, 3.0], shape=[3], name='a')
b = tf.constant([1.0, 2.0, 3.0], shape=[3], name='b')
c = a + b
# 通过log_device_placement参数来输出运行每一个运算的设备。
sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
print sess.run(c)

'''
   在没有GPU的机器上运行以上代码可以得到以下输出：
Device mapping: no known devices.

add: /job:localhost/replica:0/task:0/cpu:0
b: /job:localhost/replica:0/task:0/cpu:0
a: /job:localhost/replica:0/task:0/cpu:0
[ 2.  4.  6.]
'''

在以上代码中，TensorFlow程序生成会话时加入了参数log_device_placement=True，所以程序会将运行每一个操作的设备输出到屏幕。于是除了可以看到最后的计算结果之外，还可以看到类似“add:/job:localhost/replica:0/task:0/cpu:0”这样的输出。这些输出显示了执行每一个运算的设备。比如加法操作add是通过CPU来运行的，因为它的设备名称中包含了/cpu:0。

在配置好GPU环境的TensorFlow中，如果操作没有明确地指定运行设备，那么TensorFlow会优先选择GPU。比如将以上代码在亚马逊（Amazon Web Services, AWS）的 g2.8xlarge实例上运行时，会得到以下运行结果。

Device mapping:
/job:localhost/replica:0/task:0/gpu:0 -> device: 0, name: GRID K520, pci bus id: 0000:00:03.0
/job:localhost/replica:0/task:0/gpu:1 -> device: 1, name: GRID K520, pci bus id: 0000:00:04.0
/job:localhost/replica:0/task:0/gpu:2 -> device: 2, name: GRID K520, pci bus id: 0000:00:05.0
/job:localhost/replica:0/task:0/gpu:3 -> device: 3, name: GRID K520, pci bus id: 0000:00:06.0

add: /job:localhost/replica:0/task:0/gpu:0
b: /job:localhost/replica:0/task:0/gpu:0
a: /job:localhost/replica:0/task:0/gpu:0
[ 2.  4.  6.]

从上面的输出可以看到在配置好GPU环境的TensorFlow中，TensorFlow会自动优先将运算放置在GPU上。不过，尽管g2.8xlarge实例有4个GPU，在默认情况下，TensorFlow只会将运算优先放到/gpu:0上。于是可以看见在上面的程序中，所有的运算都被放在了/gpu:0上。如果需要将某些运算放到不同的GPU或者CPU上，就需要通过tf.device来手工指定。下面的程序给出了一个通过tf.device手工指定运行设备的样例。

import tensorflow as tf

# 通过tf.device将运算指定到特定的设备上。
with tf.device('/cpu:0'):
   a = tf.constant([1.0, 2.0, 3.0], shape=[3], name='a')
   b = tf.constant([1.0, 2.0, 3.0], shape=[3], name='b')
with tf.device('/gpu:1'):
    c = a + b

sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
print sess.run(c)

'''
在AWS g2.8xlarge实例上运行上述代码可以得到一下结果：
Device mapping:
/job:localhost/replica:0/task:0/gpu:0 -> device: 0, name: GRID K520, pci bus id: 0000:00:03.0
/job:localhost/replica:0/task:0/gpu:1 -> device: 1, name: GRID K520, pci bus id: 0000:00:04.0
/job:localhost/replica:0/task:0/gpu:2 -> device: 2, name: GRID K520, pci bus id: 0000:00:05.0
/job:localhost/replica:0/task:0/gpu:3 -> device: 3, name: GRID K520, pci bus id: 0000:00:06.0

add: /job:localhost/replica:0/task:0/gpu:1
b: /job:localhost/replica:0/task:0/cpu:0
a: /job:localhost/replica:0/task:0/cpu:0
[ 2.  4.  6.]
'''

在以上代码中可以看到生成常量a和b的操作被加载到了CPU上，而加法操作被放到了第二个GPU“/gpu:1”上。在TensorFlow中，不是所有的操作都可以被放在GPU上，如果强行将无法放在GPU上的操作指定到GPU上，那么程序将会报错。

多GPU并行

下面将给出具体的TensorFlow代码在一台机器的多个GPU上并行训练深度学习模型。因为一般来说一台机器上的多个GPU性能相似，所以在这种设置下会更多地采用同步模式训练深度学习模型。下面将给出具体的代码，在多GPU上训练深度学习模型解决MNIST问题。（该代码可以在TensorFlow 0.9.0下运行，对于更新TensorFlow的版本，请参考Github代码库：https://github.com/caicloud/tensorflow-tutorial）

# -*- coding: utf-8 -*-

from datetime import datetime
import os
import time

import tensorflow as tf

# 定义训练神经网络时需要用到的配置。
BATCH_SIZE = 100 
LEARNING_RATE_BASE = 0.001
LEARNING_RATE_DECAY = 0.99
REGULARAZTION_RATE = 0.0001
TRAINING_STEPS = 1000
MOVING_AVERAGE_DECAY = 0.99 
N_GPU = 4

# 定义日志和模型输出的路径。
MODEL_SAVE_PATH = "/path/to/logs_and_models/"
MODEL_NAME = "model.ckpt"

# 定义数据存储的路径。因为需要为不同的GPU提供不同的训练数据，所以通过placerholder
# 的方式就需要手动准备多份数据。为了方便训练数据的获取过程，可以采用输入队列的方式从
# TFRecord中读取数据。于是在这里提供的数据文件路径为将MNIST训练数据转化为
# TFRecords格式之后的路径。
DATA_PATH = "/path/to/data.tfrecords" 

# 定义输入队列得到训练数据，具体细节可以参考《TensorFlow:实战Google深度学习框架》
# 第七章。
def get_input():
    filename_queue = tf.train.string_input_producer([DATA_PATH]) 
    reader = tf.TFRecordReader()
_, serialized_example = reader.read(filename_queue)
# 定义数据解析格式。
    features = tf.parse_single_example(
        serialized_example,
        features={
            'image_raw': tf.FixedLenFeature([], tf.string),
            'pixels': tf.FixedLenFeature([], tf.int64),
            'label': tf.FixedLenFeature([], tf.int64),
        })        
    # 解析图片和标签信息。
    decoded_image = tf.decode_raw(features['image_raw'], tf.uint8)
    reshaped_image = tf.reshape(decoded_image, [784])
    retyped_image = tf.cast(reshaped_image, tf.float32)
    label = tf.cast(features['label'], tf.int32)

    # 定义输入队列并返回。
    min_after_dequeue = 10000
    capacity = min_after_dequeue + 3 * BATCH_SIZE
return tf.train.shuffle_batch(
        [retyped_image, label], 
        batch_size=BATCH_SIZE, 
        capacity=capacity, 
        min_after_dequeue=min_after_dequeue)

# 定义损失函数。对于给定的训练数据、正则化损失计算规则和命名空间，计算在这个命名空间
# 下的总损失。之所以需要给定命名空间是因为不同的GPU上计算得出的正则化损失都会加入名为
# loss的集合，如果不通过命名空间就会将不同GPU上的正则化损失都加进来。
def get_loss(x, y_, regularizer, scope):
    # 沿用第四篇文章中定义的卷积神经网络计算前向传播结果。
y = inference(x, regularizer)
# 计算交叉熵损失。
cross_entropy = tf.reduce_mean(
    tf.nn.sparse_softmax_cross_entropy_with_logits(y, y_))
# 计算当前GPU上计算得到的正则化损失。
regularization_loss = tf.add_n(tf.get_collection('losses', scope))
# 计算最终的总损失。
    loss = cross_entropy + regularization_loss
    return loss

    # 计算每一个变量梯度的平均值。
def average_gradients(tower_grads):
average_grads = []
# 枚举所有的变量和变量在不同GPU上计算得出的梯度。
for grad_and_vars in zip(*tower_grads):
          # 计算所有GPU上的梯度平均值。
          grads = []
          for g, _ in grad_and_vars:
              expanded_g = tf.expand_dims(g, 0)
              grads.append(expanded_g)
          grad = tf.concat(0, grads)
          grad = tf.reduce_mean(grad, 0)

          v = grad_and_vars[0][7]
          grad_and_var = (grad, v)
          # 将变量和它的平均梯度对应起来。
          average_grads.append(grad_and_var)
    # 返回所有变量的平均梯度，这将被用于变量更新。
    return average_grads

# 主训练过程。
def main(argv=None): 
    # 将简单的运算放在CPU上，只有神经网络的训练过程放在GPU上。
with tf.Graph().as_default(), tf.device('/cpu:0'):
        # 获取训练batch。
        x, y_ = get_input()
        regularizer = tf.contrib.layers.l2_regularizer(REGULARAZTION_RATE)

        # 定义训练轮数和指数衰减的学习率。
    global_step = tf.get_variable(
        'global_step', [], initializer=tf.constant_initializer(0), 
         trainable=False)
        learning_rate = tf.train.exponential_decay(
            LEARNING_RATE_BASE, global_step, 60000 / BATCH_SIZE, 
            LEARNING_ RATE_DECAY)       

        # 定义优化方法。
        opt = tf.train.GradientDescentOptimizer(learning_rate)

        tower_grads = []
        # 将神经网络的优化过程跑在不同的GPU上。
        for i in range(N_GPU):
            # 将优化过程指定在一个GPU上。
            with tf.device('/gpu:%d' % i):
                with tf.name_scope('GPU_%d' % i) as scope:
                    cur_loss = get_loss(x, y_, regularizer, scope)
                    # 在第一次声明变量之后，将控制变量重用的参数设置为True。这样可以
                    # 让不同的GPU更新同一组参数。注意tf.name_scope函数并不会影响
    # tf.get_ variable的命名空间。
                    tf.get_variable_scope().reuse_variables()

                    # 使用当前GPU计算所有变量的梯度。
                    grads = opt.compute_gradients(cur_loss)
                    tower_grads.append(grads)

        # 计算变量的平均梯度，并输出到TensorBoard日志中。
        grads = average_gradients(tower_grads)
        for grad, var in grads:
            if grad is not None:
                tf.histogram_summary(
                    'gradients_on_average/%s' % var.op.name, grad)

        # 使用平均梯度更新参数。
        apply_gradient_op = opt.apply_gradients(
            grads, global_step=global_ step)
        for var in tf.trainable_variables():
            tf.histogram_summary(var.op.name, var)

        # 计算变量的滑动平均值。
        variable_averages = tf.train.ExponentialMovingAverage(
            MOVING_AVERAGE_DECAY, global_step)
        variables_averages_op = variable_averages.apply(
            tf.trainable_variables())

        # 每一轮迭代需要更新变量的取值并更新变量的滑动平均值。
        train_op = tf.group(apply_gradient_op, variables_averages_op)

        saver = tf.train.Saver(tf.all_variables())
        summary_op = tf.merge_all_summaries()        
        init = tf.initialize_all_variables()

        # 训练过程。
        with tf.Session(config=tf.ConfigProto(
            allow_soft_placement=True, 
            log_device_placement=True)) as sess:
            # 初始化所有变量并启动队列。
            init.run()
            coord = tf.train.Coordinator()
            threads = tf.train.start_queue_runners(sess=sess, coord=coord)
            summary_writer = tf.train.SummaryWriter(
                MODEL_SAVE_PATH, sess.graph)

            for step in range(TRAINING_STEPS):
                # 执行神经网络训练操作，并记录训练操作的运行时间。
                start_time = time.time()
                _, loss_value = sess.run([train_op, cur_loss])
                duration = time.time() - start_time

                # 每隔一段时间展示当前的训练进度，并统计训练速度。
                if step != 0 and step % 10 == 0:
                    # 计算使用过的训练数据个数。因为在每一次运行训练操作时，每一个GPU 
                    # 都会使用一个batch的训练数据，所以总共用到的训练数据个数为
                    # batch大小×GPU个数。
                    num_examples_per_step = BATCH_SIZE * N_GPU

                    # num_examples_per_step为本次迭代使用到的训练数据个数， 
                    # duration为运行当前训练过程使用的时间，于是平均每秒可以处理的训
                    # 练数据个数为num_examples_per_step / duration。
                    examples_per_sec = num_examples_per_step / duration

                    # duration为运行当前训练过程使用的时间，因为在每一个训练过程中， 
                    # 每一个GPU都会使用一个batch的训练数据，所以在单个batch上的训
                    # 练所需要时间为duration / GPU个数。
                    sec_per_batch = duration / N_GPU

                    # 输出训练信息。
                    format_str = ('step %d, loss = %.2f (%.1f examples/ '
                                    ' sec; %.3f sec/batch)')
                    print(format_str % (step, loss_value, 
                                            examples_per_sec, sec_per_batch))

                    # 通过TensorBoard可视化训练过程。
                    summary = sess.run(summary_op)
                    summary_writer.add_summary(summary, step)

                # 每隔一段时间保存当前的模型。
                if step % 1000 == 0 or (step + 1) == TRAINING_STEPS:
                    checkpoint_path = os.path.join(
                        MODEL_SAVE_PATH, MODEL_ NAME)
                    saver.save(sess, checkpoint_path, global_step=step)

            coord.request_stop()
            coord.join(threads)

if __name__ == '__main__':
    tf.app.run()

'''
在AWS的g2.8xlarge实例上运行上面这段程序可以得到类似下面的结果：
step 10, loss = 71.90 (15292.3 examples/sec; 0.007 sec/batch)
step 20, loss = 37.97 (18758.3 examples/sec; 0.005 sec/batch)
step 30, loss = 9.54 (16313.3 examples/sec; 0.006 sec/batch)
step 40, loss = 11.84 (14199.0 examples/sec; 0.007 sec/batch)
...
step 980, loss = 0.66 (15034.7 examples/sec; 0.007 sec/batch)
step 990, loss = 1.56 (16134.1 examples/sec; 0.006 sec/batch)
'''

(点击放大图像)

图1 在AWS的g2.8xlarge实例上运行MNIST样例程序时GPU的使用情况

在AWS的g2.8xlarge实例上运行以上代码可以同时使用4个GPU训练神经网络。图1显示了运行样例代码时不同GPU的使用情况。因为运行的神经网络规模比较小，所以在图1中显示的GPU使用率不高。如果训练大型的神经网络模型，TensorFlow将会占满所有用到的GPU。

(点击放大图像)

图2 使用了4个GPU的TensorFlow计算图可视化结果

图2展示了通过TensorBoard可视化得到的样例代码TensorFlow计算图，其中节点上的颜色代表了不同的设备，比如黑色代表CPU、白色代表第一个GPU，等等。从图2中可以看出，训练神经网络的主要过程被放到了GPU_0、GPU_1、GPU_2和GPU_3这4个模块中，而且每一个模块运行在一个GPU上。

分布式TensorFlow

通过多GPU并行的方式可以达到很好的加速效果。然而一台机器上能够安装的GPU有限，要进一步提升深度学习模型的训练速度，就需要将TensorFlow分布式运行在多台机器上。以下代码展示了如何创建一个最简单的TensorFlow集群：

import tensorflow as tf
c = tf.constant("Hello, distributed TensorFlow!")
# 创建一个本地TensorFlow集群
server = tf.train.Server.create_local_server()  
# 在集群上创建一个会话。
sess = tf.Session(server.target)      
# 输出Hello, distributed TensorFlow!
print sess.run(c)　

在以上代码中，首先通过 tf.train.Server.create_local_server函数在本地建立了一个只有一台机器的TensorFlow集群。然后在该集群上生成了一个会话，并通过生成的会话将运算运行在创建的TensorFlow集群上。虽然这只是一个单机集群，但它大致反应了TensorFlow集群的工作流程。TensorFlow集群通过一系列的任务（tasks）来执行TensorFlow计算图中的运算。一般来说，不同任务跑在不同机器上。最主要的例外是使用GPU时，不同任务可以使用同一台机器上的不同GPU。TensorFlow集群中的任务也会被聚合成工作（jobs），每个工作可以包含一个或者多个任务。比如在训练深度学习模型时，一台运行反向传播的机器是一个任务，而所有运行反向传播机器的集合是一种工作。

上面的样例代码是只有一个任务的集群。当一个TensorFlow集群有多个任务时，需要使用tf.train.ClusterSpec来指定运行每一个任务的机器。比如以下代码展示了在本地运行有两个任务的TensorFlow集群。第一个任务的代码如下：

import tensorflow as tf
c = tf.constant("Hello from server1!")

# 生成一个有两个任务的集群，一个任务跑在本地2222端口，另外一个跑在本地2223端口。
cluster = tf.train.ClusterSpec(
    {"local": ["localhost:2222", "localhost: 2223"]})
# 通过上面生成的集群配置生成Server，并通过job_name和task_index指定当前所启动
# 的任务。因为该任务是第一个任务，所以task_index为0。
server = tf.train.Server(cluster, job_name="local", task_index=0)

# 通过server.target生成会话来使用TensorFlow集群中的资源。通过设置
# log_device_placement可以看到执行每一个操作的任务。
sess = tf.Session(
    server.target, config=tf.ConfigProto(log_device_placement=True)) 
print sess.run(c)
server.join()

下面给出了第二个任务的代码：

import tensorflow as tf
c = tf.constant("Hello from server2!")

# 和第一个程序一样的集群配置。集群中的每一个任务需要采用相同的配置。
cluster = tf.train.ClusterSpec(
    {"local": ["localhost:2222", "localhost: 2223"]})
# 指定task_index为1，所以这个程序将在localhost:2223启动服务。
server = tf.train.Server(cluster, job_name="local", task_index=1)
# 剩下的代码都和第一个任务的代码一致。
...

启动第一个任务后，可以得到类似下面的输出：

I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:206] Initialize HostPortsGrpcChannelCache for job local -> {localhost:2222, localhost:2223}
I tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc:202] Started server with target: grpc://localhost:2222
E1123 08:26:06.824503525   12232 tcp_client_posix.c:173]     failed to connect to 'ipv4:127.0.0.1:2223': socket error: connection refused
E1123 08:26:08.825022513   12232 tcp_client_posix.c:173]     failed to connect to 'ipv4:127.0.0.1:2223': socket error: connection refused
I tensorflow/core/common_runtime/simple_placer.cc:818] Const: /job:local/ replica:0/task:0/cpu:0
Const: /job:local/replica:0/task:0/cpu:0
Hello from server1!

从第一个任务的输出中可以看到，当只启动第一个任务时，程序会停下来等待第二个任务启动，而且持续输出failed to connect to 'ipv4:127.0.0.1:2223': socket error: connection refused。当第二个任务启动后，可以看到从第一个任务中会输出Hello from server1!的结果。第二个任务的输出如下：

I tensorflow/core/distributed_runtime/rpc/grpc_channel.cc:206] Initialize HostPortsGrpcChannelCache for job local -> {localhost:2222, localhost:2223} 
I tensorflow/core/distributed_runtime/rpc/grpc_server_lib.cc:202] Started server with target: grpc://localhost:2223 
Const: /job:local/replica:0/task:0/cpu:0 
I tensorflow/core/common_runtime/simple_placer.cc:818] Const: /job:local/ replica:0/task:0/cpu:0 
Hello from server2!

值得注意的是第二个任务中定义的计算也被放在了设备/job:local/replica:0/task:0/cpu:0上。也就是说这个计算将由第一个任务来执行。从上面这个样例可以看到，通过tf.train.Server.target生成的会话可以统一管理整个TensorFlow集群中的资源。

和使用多GPU类似，TensorFlow支持通过tf.device来指定操作运行在哪个任务上。比如将第二个任务中定义计算的语句改为以下代码，就可以看到这个计算将被调度到/job:local/replica:0/task:1/cpu:0上面。

with tf.device("/job:local/task:1"):
    c = tf.constant("Hello from server2!")

在上面的样例中只定义了一个工作“local”。但一般在训练深度学习模型时，会定义两个工作。一个工作专门负责存储、获取以及更新变量的取值，这个工作所包含的任务统称为参数服务器（parameter server，ps）。另外一个工作负责运行反向传播算法来获取参数梯度，这个工作所包含的任务统称为计算服务器（worker）。下面给出了一个比较常见的用于训练深度学习模型的TensorFlow集群配置方法。

tf.train.ClusterSpec({
    "worker": [
        "tf-worker0:2222",
        "tf-worker1:2222",
        "tf-worker2:2222"
    ],
    "ps": [
        "tf-ps0:2222",
        "tf-ps1:2222"
    ]})

使用分布式TensorFlow训练深度学习模型一般有两种方式。一种方式叫做计算图内分布式（in-graph replication）。使用这种分布式训练方式时，所有的任务都会使用一个TensorFlow计算图中的变量（也就是深度学习模型中的参数），而只是将计算部分发布到不同的计算服务器上。

上面给出的使用多GPU样例程序就是这种方式。多GPU样例程序将计算复制了多份，每一份放到一个GPU上进行运算。但不同的GPU使用的参数都是在一个TensorFlow计算图中的。因为参数都是存在同一个计算图中，所以同步更新参数比较容易控制。在上面给出的代码也实现了参数的同步更新。然而因为计算图内分布式需要有一个中心节点来生成这个计算图并分配计算任务，所以当数据量太大时，这个中心节点容易造成性能瓶颈。

另外一种分布式TensorFlow训练深度学习模型的方式叫计算图之间分布式（between-graph replication）。使用这种分布式方式时，在每一个计算服务器上都会创建一个独立的TensorFlow计算图，但不同计算图中的相同参数需要以一种固定的方式放到同一个参数服务器上。TensorFlow提供了tf.train.replica_device_setter函数来帮助完成这一个过程。

因为每个计算服务器的TensorFlow计算图是独立的，所以这种方式的并行度要更高。但在计算图之间分布式下进行参数的同步更新比较困难。

为了解决这个问题，TensorFlow提供了tf.train.SyncReplicasOptimizer函数来帮助实现参数的同步更新。这让计算图之间分布式方式被更加广泛地使用。因为篇幅所限，所以这里不在给出具体代码，具体代码可以在Github代码库中找到。

使用Caicloud TaaS平台运行分布式TensorFlow

每次运行分布式TensorFlow都需要登录不同的机器来启动集群。这使得使用起来非常不方便。当需要使用100台机器运行分布式TensorFlow时，需要手动登录到每一台机器并启动TensorFlow服务，这个过程十分繁琐。而且，当某个服务器上的程序死掉之后，TensorFlow并不能自动重启，这给监控工作带来了巨大的难度。

如果类比TensorFlow与Hadoop，可以发现TensorFlow只实现了相当于Hadoop中MapReduce的计算框架，而没有提供类似Yarn的集群管理工具以及HDFS的存储系统。为了降低分布式TensorFlow的使用门槛，才云科技(Caicloud.io)基于Kubernetes容器云平台提供了一个分布式TensorFlow平台TensorFlow as a Service（TaaS）。

从我们提供的开源代码库中可以看出，编写分布式TensorFlow程序需要指定很多与模型训练无关的代码来完成TensorFlow集群的设置工作。为了降低分布式TensorFlow的学习成本，Caicloud的TensorFlow as a Service（TaaS）平台首先对TensorFlow集群进行了更高层的封装，屏蔽了其中与模型训练无关的底层细节。

其次，TaaS平台结合了谷歌开源的容器云平台管理工具Kubernetes来实现对分布式TensorFlow任务的管理和监控，并支持通过UI设置分布式 TensorFlow任务的节点个数、是否使用GPU等信息。