CODEWITHSUNDEEP
pythontensorflowkerasdeep-learningdropout
Daniel Möller
Daniel Möller

Reputation: 86620

Tensorflow Custom Training With Phases

I need to create a custom training loop with Tensorflow / Keras (because I want to have more than one optimizer and tell which weights each optimizer should act upon).

Although this tutorial and that one too are quite clear regarding this matter, they miss a very important point: how do I predict for training phase and how do I predict for validation phase?

Suppose my model has Dropout layers, or BatchNormalization layers. They certainly work in a completely different way whether they are in training or validation.

How do I adapt these tutorials? This is a dummy example (may contain one or two pieces of pseudocode):

# Iterate over epochs.
for epoch in range(3):


    # Iterate over the batches of the dataset.
    for step, (x_batch_train, y_batch_train) in enumerate(train_dataset):
        with tf.GradientTape() as tape:

            #model with two outputs
            #IMPORTANT: must be in training phase (use dropouts, calculate batch statistics)
            logits1, logits2 = model(x_batch_train) #must be "training"

            loss_value1 = loss_fn1(y_batch_train[0], logits1)
            loss_value2 = loss_fn2(y_batch_train[1], logits2)

            grads1 = tape.gradient(loss_value1, model.trainable_weights[selection1])    
            grads2 = tape.gradient(loss_value2, model.trainable_weights[selection2])

            optimizer1.apply_gradients(zip(grads1, model.trainable_weights[selection1]))
            optimizer2.apply_gradients(zip(grads2, model.trainable_weights[selection2]))



    # Run a validation loop at the end of each epoch.
    for x_batch_val, y_batch_val in val_dataset:

        ##Important: must be validation phase
            #dropouts are off: calculate all neurons and divide value    
            #batch norms use previously calculated statistics    
        val_logits1, val_logits2 = model(x_batch_val)

        #.... do the evaluations

Upvotes: 1

Views: 202

Answers (1)

javidcf
javidcf

Reputation: 59731

I think you can just pass a training parameter when you call a tf.keras.Model, and it will be passed down to the layers:

# On training
logits1, logits2 = model(x_batch_train, training=True)
# On evaluation
val_logits1, val_logits2 = model(x_batch_val, training=False)

Upvotes: 2

Related Questions