Одиночные обновления с использованием tf.GradientTape с несколькими выходами - PullRequest
Новая
       138

Одиночные обновления с использованием tf.GradientTape с несколькими выходами

1 голос
/ 14 июля 2020

Я определил следующую модель, у которой есть два различных выхода: 

input_layer = keras.layers.Input(shape = (1, 20), name = "input_features")

# Shared layers
hidden_1 = keras.layers.Dense(32, 
                            activation = "relu", 
                            name = "LSTM_shared_l"
                            )(input_layer)

# Additional layers
hidden_2 = keras.layers.Dense(32, 
                              activation = "selu",
                              name = "Forecasting_extra_layer_1"
                              )(input_layer)

hidden_3 = keras.layers.Dense(32, 
                              activation = "selu",
                              name = "Forecasting_extra_layer_2"
                              )(hidden_2)


# Output layers
f_output = keras.layers.Dense(1, 
                              name = "F_output")(hidden_1)

rl_output = keras.layers.Dense(32, 
                               name = "RL_output")(hidden_3)

model = keras.Model(inputs = [input_layer], outputs = [f_output, rl_output])


model.summary()

, и я хотел бы обучить ее с помощью GradientTape, выполняя отдельные итерации; только с одним выходом, я бы использовал следующий код: 

with tf.GradientTape() as tape:
        
        predictions = model(inputs)
        pred_values = tf.reduce_sum(predictions, axis=1, keepdims=True)
        loss = tf.reduce_mean(loss_fn(target_pred, pred_values))
    
    grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))

Как я могу расширить его до сценария с несколькими выходами?

1 Ответ

1 голос
/ 14 июля 2020

Существует несколько стратегий, самая простая из них - вычислить убыток для обоих выходов и суммировать результаты вместе: 

predictions_1, predictions_2 = model(inputs)

predictions_1 = ...
predictions_2 = ... # any desired post-processing 

loss = tf.reduce_mean(loss_fn(target_1, predictions_1)) + tf.reduce_mean(loss_fn(target_2, predictions_2))

, а затем вы можете безопасно спуститься: 

grads = tape.gradient(loss, model.trainable_variables)
optimizer.apply_gradients(zip(grads, model.trainable_variables))
...