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

Question

为了测试，我将一个模型分成两个模型，我想计算损失并将梯度应用于两个模型，就像它是一个模型一样。

这是我的两个简单模型：

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，将它们的输出发送到第三个模型。并让每个模型 1 在两个 GPU 上进行训练：

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)

Answer 1

@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))