how to use multiple gradients with TensorFlow GradientTape?

user_14596305 0
Viewed 852

Having 3 neural networks connected like in the below code, how can we take two gradients from the initial network?? first gradient work but the second one returning the None tensor. seems like they are not related two each other to get the gradient. what is the problem here??

with tf.GradientTape() as tape1:
    with tf.GradientTape() as tape2:
        output1 = NN_model1(input1, training=True)
        output2 = NN_model2(output1, training=True)
        output3 = NN_model3([input1, output1, output2], training=True)
        loss1 = -tf.math.reduce_mean(output3)
        loss2 = -tf.math.reduce_mean(output2)
    grad1 = tape2.gradient(loss1, NN_model1.trainable_variables)
grad2 = tape1.gradient(loss2, grad1)
optimizer.apply_gradients(zip(grad2, NN_model1.trainable_variables))
keras tensorflow gradient neural-network python
1 Answers

I think correct approach should be as follows:

with tf.GradientTape() as tape:
    output1 = NN_model1(input1, training=True)
    output2 = NN_model2(output1, training=True)
    output3 = NN_model3([input1, output1, output2], training=True)
    loss1 = -tf.math.reduce_mean(output3)
    loss2 = -tf.math.reduce_mean(output2)
grad = tape.gradient([loss1, loss2], NN_model1.trainable_variables)
optimizer.apply_gradients(zip(grad, NN_model1.trainable_variables))
user_92157800
Related