Tensorflow多GPU训练模型

Thanks to : http://www.tensorfly.cn/tfdoc/api_docs/SOURCE/tutorials/deep_cnn.html

现代的工作站可能包含多个 GPU 进行科学计算。TensorFlow 可以利用这一环境在多个 GPU 卡上运行训练程序。

在并行、分布式的环境中进行训练，需要对训练程序进行协调。对于接下来的描述，术语模型拷贝（model replica）特指在一个数据子集中训练出来的模型的一份拷贝。

如果天真的对模型参数的采用异步方式更新将会导致次优的训练性能，这是因为我们可能会基于一个旧的模型参数的拷贝去训练一个模型。但与此相反采用完全同步更新的方式，其速度将会变得和最慢的模型一样慢 (Conversely, employing fully synchronous updates will be as slow as the slowest model replica.)。

在具有多个 GPU 的工作站中，每个 GPU 的速度基本接近，并且都含有足够的内存来运行整个 CIFAR-10 模型。因此我们选择以下方式来设计我们的训练系统：

在每个 GPU 上放置单独的模型副本；

等所有 GPU 处理完一批数据后再同步更新模型的参数；

示意图如下：

多GPU训练机制示意图

可以看到，每一个 GPU 会用一批独立的数据计算梯度和估计值。这种设置可以非常有效的将一大批数据分割到各个 GPU 上。

这一机制要求所有 GPU 能够共享模型参数。但是众所周知在 GPU 之间传输数据非常的慢，因此我们决定在 CPU 上存储和更新所有模型的参数 (对应图中绿色矩形的位置)。这样一来，GPU 在处理一批新的数据之前会更新一遍的参数。

图中所有的 GPU 是同步运行的。所有 GPU 中的梯度会累积并求平均值 (绿色方框部分)。模型参数会利用所有模型副本梯度的均值来更新。

在多个设备中设置变量和操作

在多个设备中设置变量和操作时需要做一些特殊的抽象。

我们首先需要把在单个模型拷贝中计算估计值和梯度的行为抽象到一个函数中。在代码中，我们称这个抽象对象为 “tower”。对于每一个 “tower” 我们都需要设置它的两个属性：

• 在一个 tower 中为所有操作设定一个唯一的名称。tf.name_scope()通过添加一个范围前缀来提供该唯一名称。比如，第一个 tower 中的所有操作都会附带一个前缀 tower_0，示例：tower_0/conv1/Conv2D；

• 在一个 tower 中运行操作的优先硬件设备。 tf.device() 提供该信息。比如，在第一个 tower 中的所有操作都位于 device('/gpu:0') 范围中，暗含的意思是这些操作应该运行在第一块 GPU 上；

为了在多个 GPU 上共享变量，所有的变量都绑定在 CPU 上，并通过 tf.get_variable()访问。可以查看 Sharing Variables 以了解如何共享变量。