gradients return None when trying to apply gradients to 2 keras models

user_14218002 0
Viewed 66

I want to combine a cnn and a transformers and apply the gradients to both models.

I create my CNN model:

cnn_model = models.Sequential([
layers.Conv1D(filters=32, kernel_size=3,
                  strides=1, padding="causal",
                  activation="relu",
                  input_shape=[None, 1024]),
layers.MaxPooling1D(pool_size=2, strides=1),
layers.Conv1D(filters=64, kernel_size=3,
                  strides=1, padding="causal",
                  activation="relu"),
layers.MaxPooling1D(pool_size=2, strides=1),
layers.Conv1D(filters=64, kernel_size=3,
                  strides=1, padding="causal",
                  activation="relu"),
layers.BatchNormalization(), ]) 
cnn_model.compile(optimizer=optimizer)

transformer = Transformer(num_layers, d_model2, num_heads, dff,
                      input_vocab_size2, target_vocab_size2, 
                      pe_input=input_vocab_size2, 
                      pe_target=target_vocab_size2,
                      rate=dropout_rate)

My optimizers and loss function are:

def loss_function(real, pred):
    mask = tf.math.logical_not(tf.math.equal(real, 0))
    loss_ = loss_object(real, pred)

    mask = tf.cast(mask, dtype=loss_.dtype)
    loss_ *= mask

    return tf.reduce_sum(loss_)/tf.reduce_sum(mask)

loss_object = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction='none')
optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=0.9, beta_2=0.98, epsilon=1e-9)

In my training step, I have:

with tf.GradientTape() as tape:
    
    cnn_prediction = cnn_model(inp, training=True)
    
    predictions, _ = transformer(cnn_prediction, tar_inp, 
                             True, 
                             enc_padding_mask, 
                             combined_mask, 
                             dec_padding_mask)
    loss = loss_function(tar_real, predictions)

gradients = tape.gradient(loss, transformer.trainable_variables)
optimizer.apply_gradients(zip(gradients, transformer.trainable_variables))

cnn_gradients = tape.gradient(loss, cnn_model.trainable_variables)
optimizer.apply_gradients(zip(cnn_gradients, cnn_model.trainable_variables))

train_loss(loss)
train_accuracy(accuracy_function(tar_real, predictions))

but when I apply gradients for the cnn model, what I get for cnn_gradients is:

[None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None]

and the error

No gradients provided for any variable: ['conv1d_10/kernel:0', 'conv1d_10/bias:0', 'conv1d_11/kernel:0', 'conv1d_11/bias:0', 'conv1d_12/kernel:0', 'conv1d_12/bias:0', 'conv1d_13/kernel:0', 'conv1d_13/bias:0', 'conv1d_14/kernel:0', 'conv1d_14/bias:0', 'conv1d_15/kernel:0', 'conv1d_15/bias:0', 'conv1d_16/kernel:0', 'conv1d_16/bias:0', 'conv1d_17/kernel:0', 'conv1d_17/bias:0', 'conv1d_18/kernel:0', 'conv1d_18/bias:0', 'conv1d_19/kernel:0', 'conv1d_19/bias:0', 'batch_normalization_3/gamma:0', 'batch_normalization_3/beta:0'].

Any ideas on how I can make this work ? What am I missing to make this work?

Thanks in advance

keras tensorflow deep-learning
1 Answers

From the documentation of the tf.GradientTape:

By default, the resources held by a GradientTape are released as soon as GradientTape.gradient() method is called. To compute multiple gradients over the same computation, create a persistent gradient tape. This allows multiple calls to the gradient() method as resources are released when the tape object is garbage collected.

You should either use two tapes, or set persistent=True when creating your tape object, or update all your gradients at the same time:

Using persistent=True

with tf.GradientTape(persistent=True) as tape:
    intermediate_result = model1(inp)
    y_pred = model2(intermediate_result)
    loss = loss_func(y_true, y_pred)
model1_gradients = tape.gradient(loss, model1.trainable_variables)
model2_gradients = tape.gradient(loss, model2.trainable_variables)

Be aware that setting persistent=True can have some impact on performance.

Using 2 tapes

with tf.GradientTape() as tape1, tf.GradientTape() as tape2:
    intermediate_result = model1(inp)
    y_pred = model2(intermediate_result)
    loss = loss_func(y_true, y_pred)
model1_gradients = tape1.gradient(loss, model1.trainable_variables)
model2_gradients = tape2.gradient(loss, model2.trainable_variables)

Calculating the gradient in one pass:

with tf.GradientTape() as tape:
    intermediate_result = model1(inp)
    y_pred = model2(intermediate_result)
    loss = loss_func(y_true, y_pred)
all_gradients = tape.gradient(loss, model1.trainable_variables + model2.trainable_variables)
# applying the gradient in one pass as well
optimizer.apply_gradients(zip(all_gradients, model1.trainable_variables + model2.trainable_variables))

I would probably favor the last option, which looks the simplest to me.

user_73701530
Related