在第11章，我们讨论了几种可以明显加速训练的技术：更好的权重初始化，批量标准化，复杂的优化器等等。 但是，即使采用了所有这些技术，在具有单个CPU的单台机器上训练大型神经网络可能需要几天甚至几周的时间。
在本章中，我们将看到如何使用TensorFlow在多个设备（CPU和GPU）上分配计算并将它们并行运行（参见图12-1）。 首先，我们会先在一台机器上的多个设备上分配计算，然后在多台机器上的多个设备上分配计算。
与其他神经网络框架相比，TensorFlow对分布式计算的支持是其主要亮点之一。 它使您可以完全控制如何跨设备和服务器分布（或复制）您的计算图，并且可以让您以灵活的方式并行和同步操作，以便您可以在各种并行方法之间进行选择。
我们来看一些最流行的方法来并行执行和训练一个神经网络，这让我们不再需要等待数周才能完成训练算法，而最终可能只会等待几个小时。 这不仅可以节省大量时间，还意味着您可以更轻松地尝试各种模型，并经常重新训练模型上的新数据。
还有其他很好的并行化例子，包括当我们在微调模型时可以探索更大的超参数空间，并有效地运行大规模神经网络。
但我们必须先学会走路才能跑步。 我们先从一台机器上的几个GPU上并行化简单图形开始。
只需添加GPU卡到单个机器，您就可以获得主要的性能提升。 事实上，在很多情况下，这就足够了。 你根本不需要使用多台机器。 例如，通常在单台机器上使用8个GPU，而不是在多台机器上使用16个GPU（由于多机器设置中的网络通信带来的额外延迟），可以同样快地训练神经网络。
在本节中，我们将介绍如何设置您的环境，以便TensorFlow可以在一台机器上使用多个GPU卡。 然后，我们将看看如何在可用设备上进行分布操作，并且并行执行它们。
为了在多个GPU卡上运行TensorFlow，首先需要确保GPU卡具有NVidia计算能力（大于或等于3.0）。 这包括Nvidia的Titan，Titan X，K20和K40（如果你拥有另一张卡，你可以在https://developer.nvidia.com/cuda-gpus查看它的兼容性）。
如果您不拥有任何GPU卡，则可以使用具有GPU功能的主机服务器，如Amazon AWS。 在ŽigaAvsec的博客文章中，提供了在Amazon AWS GPU实例上使用Python 3.5设置TensorFlow 0.9的详细说明。将它更新到最新版本的TensorFlow应该不会太难。 Google还发布了一项名为Cloud Machine Learning的云服务来运行TensorFlow图表。 2016年5月，他们宣布他们的平台现在包括配备 tensor processing units（TPU）的服务器，专门用于机器学习的处理器，比许多GPU处理ML任务要快得多。 当然，另一种选择只是购买你自己的GPU卡。 Tim Dettmers写了一篇很棒的博客文章来帮助你选择，他会定期更新它。
您必须下载并安装相应版本的CUDA和cuDNN库（如果您使用的是TensorFlow 1.0.0，则为CUDA 8.0和cuDNN 5.1），并设置一些环境变量，以便TensorFlow知道在哪里可以找到CUDA和cuDNN。 详细的安装说明可能会相当迅速地更改，因此最好按照TensorFlow网站上的说明进行操作。
Nvidia的CUDA允许开发者使用支持CUDA的GPU进行各种计算（不仅仅是图形加速）。 Nvidia的CUDA深度神经网络库（cuDNN）是针对DNN的GPU加速原语库。 它提供了常用DNN计算的优化实现，例如激活层，归一化，前向和后向卷积以及池化（参见第13章）。 它是Nvidia Deep Learning SDK的一部分（请注意，它需要创建一个Nvidia开发者帐户才能下载它）。 TensorFlow使用CUDA和cuDNN来控制GPU卡并加速计算（见图12-2）。
您可以使用nvidia-smi命令来检查CUDA是否已正确安装。 它列出了可用的GPU卡以及每张卡上运行的进程：
最后，您必须安装支持GPU的TensorFlow。 如果你使用virtualenv创建了一个独立的环境，你首先需要激活它：
$ cd $ML_PATH
# Your ML working directory (e.g., $HOME/ml) $ source env/bin/activate
然后安装合适的支持GPU的TensorFlow版本：
$ pip3 install --upgrade tensorflow-gpu
现在您可以打开一个Python shell并通过导入TensorFlow并创建一个会话来检查TensorFlow是否正确检测并使用CUDA和cuDNN：
>>> import tensorflow as tf
I [...]/dso_loader.cc:108] successfully opened CUDA library libcublas.so locally
I [...]/dso_loader.cc:108] successfully opened CUDA library libcudnn.so locally
I [...]/dso_loader.cc:108] successfully opened CUDA library libcufft.so locally
I [...]/dso_loader.cc:108] successfully opened CUDA library libcuda.so.1 locally
I [...]/dso_loader.cc:108] successfully opened CUDA library libcurand.so locally
>>> sess = tf.Session()
[...]
I [...]/gpu_init.cc:102] Found device 0 with properties:
name: GRID K520
major: 3 minor: 0 memoryClockRate (GHz) 0.797
pciBusID 0000:00:03.0
Total memory: 4.00GiB
Free memory: 3.95GiB
I [...]/gpu_init.cc:126] DMA: 0
I [...]/gpu_init.cc:136] 0: Y
I [...]/gpu_device.cc:839] Creating TensorFlow device
(/gpu:0) -> (device: 0, name: GRID K520, pci bus id: 0000:00:03.0)
看起来不错！ TensorFlow检测到CUDA和cuDNN库，并使用CUDA库来检测GPU卡（在这种情况下是Nvidia Grid K520卡）。
默认情况下，TensorFlow会在您第一次运行图形时自动获取所有可用GPU中的所有RAM，因此当第一个程序仍在运行时，您将无法启动第二个TensorFlow程序。 如果你尝试，你会得到以下错误：
E [...]/cuda_driver.cc:965] failed to allocate 3.66G (3928915968 bytes) from device: CUDA_ERROR_OUT_OF_MEMORY
一种解决方案是在不同的GPU卡上运行每个进程。 为此，最简单的选择是设置CUDA_VISIBLE_DEVICES环境变量，以便每个进程只能看到对应的GPU卡。 例如，你可以像这样启动两个程序：
$ CUDA_VISIBLE_DEVICES=0,1 python3 program_1.py
# and in another terminal:
$ CUDA_VISIBLE_DEVICES=3,2 python3 program_2.py
程序＃1只会看到GPU卡0和1（分别编号为0和1），程序＃2只会看到GPU卡2和3（分别编号为1和0）。 一切都会正常工作（见图12-3）。
另一种选择是告诉TensorFlow只抓取一小部分内存。 例如，要使TensorFlow只占用每个GPU内存的40％，您必须创建一个ConfigProto对象，将其gpu_options.per_process_gpu_memory_fraction选项设置为0.4，并使用以下配置创建session：
config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.4
session = tf.Session(config=config)
现在像这样的两个程序可以使用相同的GPU卡并行运行（但不是三个，因为3×0.4> 1）。 见图12-4。
如果在两个程序都运行时运行nvidia-smi命令，则应该看到每个进程占用每个卡的总RAM大约40％：
另一种选择是告诉TensorFlow只在需要时才抓取内存。 为此，您必须将config.gpu_options.allow_growth设置为True。 但是，TensorFlow一旦抓取内存就不会释放内存（以避免内存碎片），因此您可能会在一段时间后内存不足。 是否使用此选项可能难以确定，因此一般而言，您可能想要坚持之前的某个选项。
好的，现在你已经有了一个支持GPU的TensorFlow安装。 让我们看看如何使用它！
TensorFlow白皮书介绍了一种友好的动态布置器算法，该算法能够自动将操作分布到所有可用设备上，并考虑到以前运行图中所测量的计算时间，估算每次操作的输入和输出张量的大小， 每个设备可用的RAM，传输数据进出设备时的通信延迟，来自用户的提示和约束等等。 不幸的是，这种复杂的算法是谷歌内部的，它并没有在TensorFlow的开源版本中发布。它被排除在外的原因似乎是，由用户指定的一小部分放置规则实际上比动态放置器放置的更有效。 然而，TensorFlow团队正在努力改进它，并且最终可能会被开放。
在此之前，TensorFlow都是简单的放置，它（如其名称所示）非常基本。
无论何时运行图形，如果TensorFlow需要评估尚未放置在设备上的节点，则它会使用简单放置器将其放置在未放置的所有其他节点上。 简单放置尊重以下规则：
- 如果某个节点已经放置在图形的上一次运行中的某个设备上，则该节点将保留在该设备上。
- 否则，如果用户将一个节点固定到设备上（下面介绍），则放置器将其放置在该设备上。
- 否则，它默认为GPU＃0，如果没有GPU，则默认为CPU。
正如您所看到的，将操作放在适当的设备上主要取决于您。 如果您不做任何事情，整个图表将被放置在默认设备上。 要将节点固定到设备上，您必须使用device()函数创建一个设备块。 例如，以下代码将变量a和常量b固定在CPU上，但乘法节点c不固定在任何设备上，因此将放置在默认设备上：
with tf.device("/cpu:0"):
a = tf.Variable(3.0)
b = tf.constant(4.0)
c = a * b
其中，“/ cpu：0”设备合计多CPU系统上的所有CPU。 目前没有办法在特定CPU上固定节点或仅使用所有CPU的子集。
让我们检查一下简单的放置器是否遵守我们刚刚定义的布局约束条件。 为此，您可以将log_device_placement选项设置为True; 这告诉放置器在放置节点时记录消息。 例如：
>>> config = tf.ConfigProto()
>>> config.log_device_placement = True
>>> sess = tf.Session(config=config)
I [...] Creating TensorFlow device (/gpu:0) -> (device: 0, name: GRID K520, pci bus id: 0000:00:03.0)
[...]
>>> x.initializer.run(session=sess)
I [...] a: /job:localhost/replica:0/task:0/cpu:0
I [...] a/read: /job:localhost/replica:0/task:0/cpu:0
I [...] mul: /job:localhost/replica:0/task:0/gpu:0
I [...] a/Assign: /job:localhost/replica:0/task:0/cpu:0
I [...] b: /job:localhost/replica:0/task:0/cpu:0
I [...] a/initial_value: /job:localhost/replica:0/task:0/cpu:0
>>> sess.run(c)
12
Info中以大写字母“I”开头的行是日志消息。 当我们创建一个会话时，TensorFlow会记录一条消息，告诉我们它已经找到了一个GPU卡（在这个例子中是Grid K520卡）。 然后，我们第一次运行图形（在这种情况下，当初始化变量a时），简单布局器运行，并将每个节点放置在分配给它的设备上。正如预期的那样，日志消息显示所有节点都放在“/ cpu：0”上，除了乘法节点，它以默认设备“/ gpu：0”结束（您可以先忽略前缀：/job:localhost /replica:0 /task:0; 我们将在一会儿讨论它）。 注意，我们第二次运行图（计算c）时，由于TensorFlow需要计算的所有节点c都已经放置，所以不使用布局器。
创建设备块时，可以指定一个函数，而不是设备名称。 TensorFlow会调用这个函数来进行每个需要放置在设备块中的操作，并且该函数必须返回设备的名称来固定操作。 例如，以下代码将固定所有变量节点到“/ cpu：0”（在本例中只是变量a）和所有其他节点到“/ gpu：0”：
def variables_on_cpu(op):
if op.type == "Variable":
return "/cpu:0"
else:
return "/gpu:0"
with tf.device(variables_on_cpu):
a = tf.Variable(3.0)
b = tf.constant(4.0)
c = a * b
您可以轻松实现更复杂的算法，例如以循环方式用GPU锁定变量。
对于在设备上运行的TensorFlow操作，它需要具有该设备的实现；这被称为内核。 许多操作对于CPU和GPU都有内核，但并非全部都是。 例如，TensorFlow没有用于整数变量的GPU内核，因此当TensorFlow尝试将变量i放置到GPU＃0时，以下代码将失败：
>>> with tf.device("/gpu:0"):
... i = tf.Variable(3)
[...]
>>> sess.run(i.initializer)
Traceback (most recent call last):
[...]
tensorflow.python.framework.errors.InvalidArgumentError: Cannot assign a device to node 'Variable': Could not satisfy explicit device specification
请注意，TensorFlow推断变量必须是int32类型，因为初始化值是一个整数。 如果将初始化值更改为3.0而不是3，或者如果在创建变量时显式设置dtype = tf.float32，则一切正常。
默认情况下，如果您尝试在操作没有内核的设备上固定操作，则当TensorFlow尝试将操作放置在设备上时，您会看到前面显示的异常。 如果您更喜欢TensorFlow回退到CPU，则可以将allow_soft_placement配置选项设置为True：
with tf.device("/gpu:0"):
i = tf.Variable(3)
config = tf.ConfigProto()
config.allow_soft_placement = True
sess = tf.Session(config=config)
sess.run(i.initializer) # the placer runs and falls back to /cpu:0
到目前为止，我们已经讨论了如何在不同设备上放置节点。 现在让我们看看TensorFlow如何并行运行这些节点。
当TensorFlow运行图时，它首先找出需要评估的节点列表，然后计算每个节点有多少依赖关系。 然后TensorFlow开始评估具有零依赖关系的节点（即源节点）。 如果这些节点被放置在不同的设备上，它们显然会被并行评估。 如果它们放在同一个设备上，它们将在不同的线程中进行评估，因此它们也可以并行运行（在单独的GPU线程或CPU内核中）。
TensorFlow管理每个设备上的线程池以并行化操作（参见图12-5）。 这些被称为 inter-op 线程池。 有些操作具有多线程内核：它们可以使用其他线程池（每个设备一个）称为 intra-op 线程池（下面写成内部线程池）。
例如，在图12-5中，操作A，B和C是源操作，因此可以立即进行评估。 操作A和B放置在GPU＃0上，因此它们被发送到该设备的内部线程池，并立即进行并行评估。 操作A正好有一个多线程内核; 它的计算被分成三部分，这些部分由内部线程池并行执行。 操作C转到GPU＃1的内部线程池。
一旦操作C完成，操作D和E的依赖性计数器将递减并且都将达到0，因此这两个操作将被发送到操作内线程池以执行。
您可以通过设置inter_op_parallelism_threads选项来控制内部线程池的线程数。 请注意，您开始的第一个会话将创建内部线程池。 除非您将use_per_session_threads选项设置为True，否则所有其他会话都将重用它们。 您可以通过设置intra_op_parallelism_threads选项来控制每个内部线程池的线程数。
在某些情况下，即使所有依赖的操作都已执行，推迟对操作的评估可能也是明智之举。例如，如果它使用大量内存，但在图形中只需要更多内存，则最好在最后一刻对其进行评估，以避免不必要地占用其他操作可能需要的RAM。 另一个例子是依赖位于设备外部的数据的一组操作。 如果它们全部同时运行，它们可能会使设备的通信带宽达到饱和，并最终导致所有等待I / O。 其他需要传递数据的操作也将被阻止。 顺序执行这些通信繁重的操作将是比较好的，这样允许设备并行执行其他操作。
推迟对某些节点的评估，一个简单的解决方案是添加控制依赖关系。 例如，下面的代码告诉TensorFlow仅在评估完a和b之后才评估x和y：
a = tf.constant(1.0)
b = a + 2.0
with tf.control_dependencies([a, b]):
x = tf.constant(3.0)
y = tf.constant(4.0)
z = x + y
显然，由于z依赖于x和y，所以评估z也意味着等待a和b进行评估，即使它并未显式存在于control_dependencies()块中。 此外，由于b依赖于a，所以我们可以通过在[b]而不是[a，b]上创建控制依赖关系来简化前面的代码，但在某些情况下，“显式比隐式更好”。
很好！现在你知道了：
- 如何以任何您喜欢的方式在多个设备上进行操作
- 这些操作如何并行执行
- 如何创建控制依赖性来优化并行执行
是时候将计算分布在多个服务器上了！
要跨多台服务器运行图形，首先需要定义一个集群。 一个集群由一个或多个TensorFlow服务器组成，称为任务，通常分布在多台机器上（见图12-6）。 每项任务都属于一项job（下面称作业）。 作业只是一组通常具有共同作用的任务，例如跟踪模型参数（例如，参数服务器parameter server通常命名为“ps”）或执行计算（这样的作业通常被命名为“ worker”）。
以下集群规范定义了两个作业“ps”和“worker”，分别包含一个任务和两个任务。 在这个例子中，机器A托管着两个TensorFlow服务器（即任务），监听不同的端口：一个是“ps”作业的一部分，另一个是“worker”作业的一部分。 机器B仅托管一台TensorFlow服务器，这是“worker”作业的一部分。
cluster_spec = tf.train.ClusterSpec({
"ps": [
"machine-a.example.com:2221", # /job:ps/task:0
],
"worker": [
"machine-a.example.com:2222", # /job:worker/task:0
"machine-b.example.com:2222", # /job:worker/task:1
]})
要启动TensorFlow服务器，您必须创建一个服务器对象，并向其传递集群规范（以便它可以与其他服务器通信）以及它自己的作业名称和任务编号。 例如，要启动第一个辅助任务，您需要在机器A上运行以下代码：
server = tf.train.Server(cluster_spec, job_name="worker", task_index=0)
每台机器只运行一个任务通常比较简单，但前面的例子表明TensorFlow允许您在同一台机器上运行多个任务（如果需要的话）。 如果您在一台机器上安装了多台服务器，则需要确保它们不会全部尝试抓取每个GPU的所有RAM，如前所述。 例如，在图12-6中，“ps”任务没有看到GPU设备，想必其进程是使用CUDA_VISIBLE_DEVICES =“ ”启动的。 请注意，CPU由位于同一台计算机上的所有任务共享。
如果您希望进程除了运行TensorFlow服务器之外什么都不做，您可以通过告诉它等待服务器使用join()方法来完成，从而阻塞主线程（否则服务器将在您的主线程退出）。 由于目前没有办法阻止服务器，这实际上会永远阻止：
server.join() # blocks until the server stops (i.e., never)
一旦所有任务启动并运行（但还什么都没做），您可以从位于任何机器上的任何进程（甚至是运行中的进程）中的客户机上的任何服务器上打开会话，并使用该会话像普通的本地会议一样。比如：
a = tf.constant(1.0)
b = a + 2
c = a * 3
with tf.Session("grpc://machine-b.example.com:2222") as sess:
print(c.eval()) # 9.0
这个客户端代码首先创建一个简单的图形，然后在位于机器B（我们称之为主机）上的TensorFlow服务器上打开一个会话，并指示它评估c。 主设备首先将操作放在适当的设备上。 在这个例子中，因为我们没有在任何设备上进行任何操作，所以主设备只将它们全部放在它自己的默认设备上 - 在这种情况下是机器B的GPU设备。 然后它只是按照客户的指示评估c，并返回结果。
客户端使用gRPC协议（Google Remote Procedure Call）与服务器进行通信。 这是一个高效的开源框架，可以调用远程函数，并通过各种平台和语言获取它们的输出。它基于HTTP2，打开一个连接并在整个会话期间保持打开状态，一旦建立连接就可以进行高效的双向通信。
数据以协议缓冲区的形式传输，这是另一种开源Google技术。 这是一种轻量级的二进制数据交换格式。
TensorFlow集群中的所有服务器都可能与集群中的任何其他服务器通信，因此请确保在防火墙上打开适当的端口。
每台TensorFlow服务器都提供两种服务：主服务和辅助服务。 主服务允许客户打开会话并使用它们来运行图形。 它协调跨任务的计算，依靠辅助服务实际执行其他任务的计算并获得结果。
通过指定作业名称，任务索引，设备类型和设备索引，可以使用设备块来锁定由任何任务管理的任何设备上的操作。 例如，以下代码将a固定在“ps”作业（即机器A上的CPU）中第一个任务的CPU，并将b固定在“worker”作业的第一个任务管理的第二个GPU （这是A机上的GPU＃1）。 最后，c没有固定在任何设备上，所以主设备将它放在它自己的默认设备上（机器B的GPU＃0设备）。
with tf.device("/job:ps/task:0/cpu:0")
a = tf.constant(1.0)
with tf.device("/job:worker/task:0/gpu:1")
b = a + 2
c = a + b
如前所述，如果您省略设备类型和索引，则TensorFlow将默认为该任务的默认设备; 例如，将操作固定到“/ job：ps / task：0”会将其放置在“ps”作业（机器A的CPU）的第一个任务的默认设备上。 如果您还省略了任务索引（例如，“/ job：ps”），则TensorFlow默认为“/ task：0”。 如果省略作业名称和任务索引，则TensorFlow默认为会话的主任务。
正如我们很快会看到的那样，在分布式设置上训练神经网络时，常见模式是将模型参数存储在一组参数服务器上（即“ps”作业中的任务），而其他任务则集中在计算上（即 ，“worker”工作中的任务）。 对于具有数百万参数的大型模型，在多个参数服务器上分割这些参数非常有用，可以降低饱和单个参数服务器网卡的风险。 如果您要将每个变量手动固定到不同的参数服务器，那将非常繁琐。 幸运的是，TensorFlow提供了replica_device_setter()函数，它以循环方式在所有“ps”任务中分配变量。 例如，以下代码将五个变量引入两个参数服务器：
with tf.device(tf.train.replica_device_setter(ps_tasks=2):
v1 = tf.Variable(1.0) # pinned to /job:ps/task:0
v2 = tf.Variable(2.0) # pinned to /job:ps/task:1
v3 = tf.Variable(3.0) # pinned to /job:ps/task:0
v4 = tf.Variable(4.0) # pinned to /job:ps/task:1
v5 = tf.Variable(5.0) # pinned to /job:ps/task:0
您不必传递ps_tasks的数量，您可以传递集群spec spec = cluster_spec，TensorFlow将简单计算“ps”作业中的任务数。
如果您在块中创建其他操作，则不仅仅是变量，TensorFlow会自动将它们连接到“/ job：worker”，默认为第一个由“worker”作业中第一个任务管理的设备。 您可以通过设置worker_device参数将它们固定到其他设备，但更好的方法是使用嵌入式设备块。 内部设备块可以覆盖在外部块中定义的作业，任务或设备。 例如：
with tf.device(tf.train.replica_device_setter(ps_tasks=2)):
v1 = tf.Variable(1.0) # pinned to /job:ps/task:0 (+ defaults to /cpu:0)
v2 = tf.Variable(2.0) # pinned to /job:ps/task:1 (+ defaults to /cpu:0)
v3 = tf.Variable(3.0) # pinned to /job:ps/task:0 (+ defaults to /cpu:0)
[...]
s = v1 + v2 # pinned to /job:worker (+ defaults to task:0/gpu:0)
with tf.device("/gpu:1"):
p1 = 2 * s # pinned to /job:worker/gpu:1 (+ defaults to /task:0)
with tf.device("/task:1"):
p2 = 3 * s # pinned to /job:worker/task:1/gpu:1
这个例子假设参数服务器是纯CPU的，这通常是这种情况，因为它们只需要存储和传送参数，而不是执行密集计算。