python-3.x - Keras 中的模型并行性

Question

我正在尝试在 Keras 中实现模型并行性。

我正在使用 Keras-2.2.4 Tensorflor-1.13.1

我的代码的粗略结构是： import tensorflow as tf import keras

def model_definition():
input0 = Input(shape = (None, None))
input1 = Input(shape = (None, None))
with tf.Session(config=tf.ConfigProto(allow_soft_placement=False, log_device_placement=True)):
model = get_some_CNN_model()
with tf.device(tf.DeviceSpec(device_type="GPU", device_index=0)):
op0 = model(input0)
with tf.device(tf.DeviceSpec(device_type="GPU", device_index=1)):
op1 = model(input1)
with tf.device(tf.DeviceSpec(device_type="CPU", device_index=0)):
concatenated_ops = concatenate([op0, op1],axis=-1, name = 'check_conc1')
mixmodel = Model(inputs=[input0, input1], outputs = concatenated_ops)
return mixmodel

mymodel = model_definition()
mymodel.fit_generator()

预期结果：训练时，op0 和 op1 的计算应分别在 gpu0 和 gpu1 上完成。

问题 1：当我有 2 个 gpus 可用时，训练工作正常。nvidia-smi 显示两个 gpus 都在使用。虽然我不确定两个 gpus 是否都在做他们预期的工作。如何确认？因为，即使我将 log_device_placement 设置为 True，我也没有看到分配给 gpu 1 的任何任务

问题 2：当我在一台有 1 个 GPU 可用的机器上运行此代码时，它仍然运行良好。预计会显示错误，因为 GPU 1 不可用。

此处显示的示例按预期工作正常。它没有显示问题 2，即在单个 GPU 上，它会引发错误。

所以我认为 keras 内部正在发生一些操纵。

我也尝试过使用 import tensorflow.python.keras 而不是 import keras，以防引起任何冲突。然而，这两个问题仍然存在。

将不胜感激有关此问题的任何线索。谢谢你。

标签： python-3.xtensorflowkerasmodel