如何在异步Actor批判方法中实现Tensorflow模型并行性?

时间:2018-08-08 23:11:30

标签: tensorflow deep-learning reinforcement-learning

我正在使用具有100个并行线程的A3C体系结构。我有2张Nvidia 1080ti卡。但是我的初始模型运行在单个GPU上。通过一些扩展,我的新模型太重了,无法在单个卡上运行。由于A3C方法是无异步模型的RL方法,因此我使用模型并行性。这意味着我使用 tf.device()在两个GPU中分配了不同的操作。但是我观察到它非常慢,并且大多数情况下操作失败。

  • 是否存在一种在不更改A3C原始体系结构的情况下实现模型并行性的理想方法?

1 个答案:

答案 0 :(得分:1)

使用多个GPU的理想方法是将整个图形的实例分别展开到每个GPU。

现在,不要运行常规的优化器minimize op。使用tf.gradients(),然后在所有GPU上使用tf.reduce_mean(),计算每个卡的梯度,然后将平均梯度应用到使用optimizer.apply_gradients()

根据您的网络规模,在多个GPU之间同步梯度更新步骤将非常昂贵,您可以执行多次前向后向传播,收集梯度并在{ {1}}步骤。

Here's Tensorflow的使用多GPU的官方代码。

由于拥有多个角色,因此您还可以将前50个角色分配给GPU1,并将后50个角色分配给GPU2。然后,您可以在几次迭代后求出梯度的平均值。

我还建议在CPU和GPU之间使用队列,以最小化RAM和VRAM之间的k成本。 Here's是Nvidia的一个很好的存储库(和论文),用于在GPU上相当强大地实现A3C。

或者考虑使用Impala之类的更新的演员评论器体系结构,该体系结构显示出资源和时间效率更高。