跨多个GPU分发Keras模型

Question

我正在尝试创建一个非常大的Keras模型，并将其分布在多个GPU上。需要说明的是，我并不是想把同一模型的多个副本放在多个GPU上；我是想把一个大模型放在多个GPU上。我一直在使用Keras中的multi_gpu_model函数，但基于我在执行此操作时遇到的大量内存不足错误，它似乎只是复制了模型，而不是像我希望的那样分发它。

我研究了Horovod，但是因为我有很多windows特定的日志工具在运行，所以我在犹豫是否要使用它。

这似乎只剩下tf.estimators可供我使用了。但是，从文档中还不清楚我将如何使用这些估计器来完成我想要做的事情。例如，tf.contrib.distribute中的哪种分发策略可以让我以我想要的方式有效地批量处理模型？

我试图用估计器做的事情是可能的吗?如果可能，我应该使用哪种策略？

Answer 1

您可以使用Estimator API。使用tf.keras.estimator.model_to_estimator转换模型

session_config = tf.ConfigProto(allow_soft_placement=True)
distribute = tf.contrib.distribute.MirroredStrategy(num_gpus=4)
run_config = tf.estimator.RunConfig(train_distribute=distribute)
your_network = tf.keras.estimator.model_to_estimator(model_fn=your_keras_model, config=run_config)
your_network.train(input_fn)

别忘了编译模型

Answer 2

您可以使用TensorFlow后端将Keras模型的不同部分手动分配给不同的GPU。This guide提供了详细的示例，this article解释了如何在TensorFlow中使用Keras。

import tensorflow as tf

with tf.device("/device:GPU:0"):
    #Create first part of your neural network

with tf.device("/device:GPU:1"):
    #Create second part of your neural network

#...

with tf.device("/device:GPU:n"):
    #Create nth part of your neural network

注意: CPU和多个GPU之间的通信延迟可能会增加训练的大量开销。

Answer 3

您需要设备并行性。Keras常见问题解答的This section提供了一个使用Keras执行此操作的示例：

# Model where a shared LSTM is used to encode two different sequences in parallel
input_a = keras.Input(shape=(140, 256))
input_b = keras.Input(shape=(140, 256))

shared_lstm = keras.layers.LSTM(64)

# Process the first sequence on one GPU
with tf.device_scope('/gpu:0'):
    encoded_a = shared_lstm(tweet_a)
# Process the next sequence on another GPU
with tf.device_scope('/gpu:1'):
    encoded_b = shared_lstm(tweet_b)

# Concatenate results on CPU
with tf.device_scope('/cpu:0'):
    merged_vector = keras.layers.concatenate([encoded_a, encoded_b],
                                             axis=-1)