单机多 GPU 加速 TensorFlow 程序

​​模型并行意味着需要根据不同模型设计不同的并行方式,并将模型中不同计算节点放在不同硬件资源上运算。相比模型并行而言,数据并行比较通用,并且能简便地实现大规模并行。

数据并行的通用性表现在它几乎适用于所有深度学习模型。我们可以 利用多块 GPU 同时进行模型训练,运行在每块 GPU上的模型会获得属于自 己的一份 batch 数据,并且这些模型都是同一个神经网络,它们共享一份模型参数。

首先导入一些相关的库以及定义一些变量。由于我们的重点不是为了 训练网络,所以训练迭代的轮数可以相应取小一点。为了能可视化计算图,还要定义 TensorBoard 日志文件的存放路径。

单机多 GPU 加速 TensorFlow 程序

定义一个函数实现网络前向传播的过程。由于这个过程在之前的章节 中己经重复了多次,所以这里对代码从简。在前向传播的过程中,还对权重参数加入了 L2 正则化。这个函数的代码如下:

单机多 GPU 加速 TensorFlow 程序

接着,定义一个函数从 TFReocrd 文件中读取原始训练样本数据并通过随机打乱汇总为一个 batch 的形式, 最后返回的就是一个 batch 的训练样本数据及其对应的 label。函数用到的 TFReocrd 文件就是在之前生成过的文件。这个函数的代码如下:

单机多 GPU 加速 TensorFlow 程序

也可以像往常一样通过 placeholder 的方式为网络提供输入数据,但是由于这里使用了多个 GPU,所以在使用这种方式时一般要通过多个placeholder 分别对应到每 一个 GPU 所使用的mini-batch 数据 。相比placeholder 的方式,这种从文件读取并打乱的方式显得更加简洁。然后,定义一个计算损失的函数,要计算的损失包括 L2 正则化损失和 交叉情损失。在计算正则化损失时使用 get_collection()函数获取当前这个 GPU 上的 L2 loss (通过 scope 限定了范围),再使用 add_n()函数将所有损 失叠加到一起得到总计的正则化损失,最后作为结果返回的是总计的正则 化损失与交叉;脑损失加和的值。这个函数的代码如下:

单机多 GPU 加速 TensorFlow 程序

我们像之前一样定义一个默认的计算图,计算图内先是获得输入的数 据,然后选择使用的优化器为随机梯度下降优化器 GradientDescentOptimizer。在使用这个优化器时,学习率一般需要自行指定,这里指定学 习率为使用往ain.exponential_decay()函数创建的指数衰减学习率。

单机多 GPU 加速 TensorFlow 程序

计算图最后的一部分是会话。在会话中,先进行所有变量 的初始化 以及队列的启动,然后在一个 for 循环中进行网络模型的迭代。每隔 100 轮 选代要展示当前的训练进度和 loss 值,并通过 summary_op 执行所有的汇总 操作。在定义会话时,使用了参数 allow_soft_placement=True。对于一些特殊 情况,例如机器上包含的 GPU与定义的n GPU数量不同,此时会话仍能继 续运行,只是将相应的计算放到了其他设备(如 CPU)进行预算。如果没 有指定这个参数并且 n_GPU 定义得偏大,那么程序在运行过程中将会发生错误。

关注小鲸融创,一起深度学习金融科技!

单机多 GPU 加速 TensorFlow 程序​​​​

相关推荐