python - 如何在 Tensorflow 中使用 keras 在 2 个 GPU 之间拆分模型？

Question

本质上，我正在寻找类似with tf.device('/device:GPU:0'keras 的东西。我想将我的操作放在不同的 GPU 上。我正在使用遵循以下行的顺序模型

...
model.compile(optimizer='adam', 
                  loss='sparse_categorical_crossentropy',
                  metrics=['accuracy'])
...

model.fit(train_images, train_labels, epochs=5)

标签： pythontensorflowkeras

Answer 1

可能您可以with K.tf.device('/gpu:1'):用作上下文。或者，如果后端是 tensorflow，那么您在 tf 中分配 gpu 的方式也应该适用于 keras。