对具有字典输入(由多个模型组成)的tf.keras神经网络使用GradientTape

import tensorflow as tf

input_data = {'input1': input1_data, 'input2': input2_data, ...}
target_data = target_data

model = tf.keras.models.Sequential([...])

loss_fn = tf.keras.losses.MeanSquaredError()

optimizer = tf.keras.optimizers.Adam(learning_rate=0.001)

@tf.function
def train_step(inputs, targets):
    with tf.GradientTape() as tape:
        predictions = model(inputs)
        loss = loss_fn(targets, predictions)
    gradients = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
    return loss

epochs = 10
batch_size = 32

for epoch in range(epochs):
    dataset = tf.data.Dataset.from_tensor_slices((input_data, target_data)).batch(batch_size)
    total_loss = 0.0
    
    for inputs, targets in dataset:
        loss = train_step(inputs, targets)
        total_loss += loss
    
    average_loss = total_loss / len(dataset)
    print('Epoch {}: loss = {}'.format(epoch+1, average_loss))