Enable dropout and disable BatchNormalization at inference

Question

I would like to enable dropout at training and inference time using Tensorflow 2.5. For this aim, I set inside my model the dropout layers with the parameter training = True.

layer = tf.keras.layers.Dropout(0.2, training = True)

Then I trained my model, and made a prediction using the following code:

prediction = model(X_test, training = False)

I deliberately put training = False in the model prediction function (the model call function) because I am also using BatchNormalization layers and I don't want they are enable during inference (unlike the dropout layers). However, I don't know if putting training = False in the model prediction function will additionally put the dropout layers to training = False (override it). Could you tell if dropout still enable at inference when I am using model(X_test, training = False) ? If not, does there is a way to enable dropout and disable BatchNormalization during inference (at same time) ?

Answer 1

During the predict, dropout is disabled in Tensorflow, it can be only enabled during training by setting parameter training=True. To enable the dropout during inference, there is no direct way as of now. But doing bit of workaround as mentioned in the below links you can achieve. 1 & 2.

Batch Normalization does not perform the same computation during training and after training, it uses batch statsistics during training and the final statistics after training(i.e the final values of the moving averages). But, if you want to change this behavior, you need to have custom layer and have training flag set as per your desire. Below is the sample snippet.

class BatchNormalization(tf.keras.layers.Layer):
   [...]
   def call(self, inputs, training=None):
       [...]