Python 跨多个GPU分布Keras模型

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

我查看了Horowod，但是因为我有很多windows特定的日志记录工具在运行，所以我对使用它犹豫不决

这似乎只剩下tf.estimators供我使用。从文档中还不清楚我将如何使用这些估计器来完成我要做的事情。例如，您选择了哪种分销策略 tf.contrib.distribute能让我以我期待的方式有效地批量生产模型吗

---

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

您可以使用TensorFlow后端手动将Keras模型的不同部分分配给不同的GPU。提供详细的示例，并解释如何将Keras与TensorFlow结合使用

将tensorflow导入为tf
使用tf.device（“/device:GPU:0”）：
#创建神经网络的第一部分
使用tf.device（“/device:GPU:1”）：
#创建神经网络的第二部分
#...
使用tf.device（“/device:GPU:n”）：
#创建神经网络的第n部分

---

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

您可以使用Estimator API。使用

tf.keras.estimator.model\u转换您的模型到\u 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)

---

不要忘记编译模型，因为您需要设备并行性。Keras常见问题解答提供了如何使用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)

---

谢谢，我决定使用它，因为它是可伸缩的，可以设计为自动调整到数字。你知道一种简单的方法来生成输入吗？因为我现在就在这里。你可以使用tf.data.Dataset。首先创建dataset对象，然后应用map func并批处理数据集。在ens中，返回dataset对象并在Estimator中使用它。您不需要指定迭代器。更多关于这个