如何使用Keras计算链式模型的损失?

时间:2020-11-05 14:43:43

标签: python tensorflow keras

为了进行测试,我将一个模型分为两个模型,我想计算损耗并将梯度应用于两个模型,就像将一个模型一样。

这是我的两个简单模型:

model1 = tf.keras.models.Sequential([
    tf.keras.layers.Dense(10, activation="relu", input_shape=(10,)),
])

model2 = tf.keras.models.Sequential([
    tf.keras.layers.Dense(10, activation="softmax", input_shape=(10,)),
])

然后我对这两个模型进行前向遍历,计算第二个模型的损耗并应用渐变:

optimizer = tf.keras.optimizers.SGD()
loss = tf.keras.losses.CategoricalCrossentropy(from_logits=True)

x = tf.random.normal((1, 10)) # Input of the 1st model
y = tf.random.normal((1, 10)) # Expected output of the 2nd model

with tf.GradientTape() as tape:
    pred1 = model1(x, training=True)
    pred2 = model2(pred1, training=True)
    loss_value2 = loss(y, pred2) # Compute the loss for the second model prediction

grads = tape.gradient(loss_value2, model2.trainable_variables)
optimizer.apply_gradients(zip(grads, model2.trainable_variables))

但是我如何获得第二个模型的第一个模型的预期输出来计算损耗并对其应用渐变呢?

编辑:

测试的最终目标是拥有两个模型1,并将它们的输出发送到单个第三个模型。并在两个GPU上训练每个模型1:

with tf.device('/gpu:0'):
    pred1_1 = model1_1(x, training=True)

with tf.device('/gpu:1'):
    pred1_2 = model1_2(x, training=True)

pred1 = tf.keras.layers.concatenate([pred1_1, pred1_2])

with tf.device('/gpu:0'):
    pred2 = model2(pred1, training=True)

1 个答案:

答案 0 :(得分:2)

@Begoodpy, 我建议您将2个模型合并为一个模型,然后像往常一样训练它。

supermodel = keras.Sequential(
    [
      model1(),
      model2(),
    ]

如果您需要对模型进行更多控制,请尝试以下操作:

all_vars = model1.trainable_variables + model2.trainable_variables

grads = tape.gradient(loss_value2, all_vars)
optimizer.apply_gradients(zip(grads, all_vars))