python - 是否可以在 tensorflow 中跨多个 GPU 拆分网络？

Question

我计划运行一个非常大的循环网络(例如 2048x5)，是否可以在 tensorflow 中的一个 GPU 上定义一层？我应该如何实现模型才能达到最佳效率。我了解 GPU 间或 GPU-CPU-GPU 通信会产生开销。

Answer 1

在 TensorFlow 中，跨多个 GPU 拆分大型模型当然是可能的，但最佳地做到这一点是一个艰巨的研究问题。通常，您需要执行以下操作:

1. 用 with tf.device(...): 包裹代码的大块连续区域 block ，命名不同的 GPU:

   with tf.device("/gpu:0"):
     # Define first layer.
   
   with tf.device("/gpu:1"):
     # Define second layer.
   
   # Define other layers, etc.
   

2. 构建优化器时，将可选参数 colocate_gradients_with_ops=True 传递给 optimizer.minimize()方法:

   loss = ...
   optimizer = tf.train.AdaGradOptimizer(0.01)
   train_op = optimizer.minimize(loss, colocate_gradients_with_ops=True)
   

3. (可选。)当您创建 tf.Session 时，您可能需要在 tf.ConfigProto 中启用“软放置” ，如果模型中的任何操作无法在 GPU 上运行:

   config = tf.ConfigProto(allow_soft_placement=True)
   sess = tf.Session(config=config)