Is there a way to activate dropout during inference in Keras, while freezing batch norm layers

Question

I am trying to use the dropout layers in my model during inference time to measure the model uncertainty as described in the method outlined by Yurin Gal

A solution is described in this post: How to calculate prediction uncertainty using Keras?, which defines a new Keras function self.f = K.function([self.graph.layers[0].input, K.learning_phase()], [self.graph.layers[-1].output])

However, this method does not hold if the used model has batch normalisation layers. As this will make the model not use the mean and variance learned during training, but set new ones depending on the current batch.

Hence, I am looking for a way to put the batch layers training parameter to false but keep the dropout layer in training mode?

I am using the Keras efficientNet B0 as model, trained on custom data keras_efficientNet

I already tried to change the layers settings myself

`
        for layer in self.graph.layers[4].layers:
            if 'batch_norm' in layer.name:
                layer._trainable = False
                layer._inbound_nodes[0].output_tensors[0]._uses_learning_phase = False
                layer._inbound_nodes[0].input_tensors[0]._uses_learning_phase = False
            if 'dropout' in layer.name:
                layer._inbound_nodes[0].output_tensors[0]._uses_learning_phase = True
            for weight in self.graph.layers[4].weights:
                if 'batch_norm' in weight.name:
                    weight._trainable = False`

Nonetheless, none of this worked.

Pedrolarben · Answer 1 · 2020-06-12T16:35:25.473

This question is duplicated, it is already answered here How to apply Monte Carlo Dropout, in tensorflow, for an LSTM if batch normalization is part of the model?

Basically, when you define your model, you should add training=True to your Dropout layer

inputs = tf.keras.Input(...) 
x = tf.keras.layers.___(...)(input)
...
x = tf.keras.layers.Dropout(...)(x, training=True)
...

In the case, you cannot modify the model constructor code, you can modify it like this (not my favourite solution) [1].

# load model
model_config = model.get_config()
layer_index = 3 # layer index you want to modify
model_config['layers'][layer_index]['inbound_nodes'][0][0][-1]['training'] = True
model = tf.keras.models.model_from_config(model_config)

thanks for the quick reply. the issue is that I am not defining the model again. I am only using it for inference and don't want to train it again.I use the load function from keras.model to load the already trained .h5 model. is there a way to change the definition of the model when loading it this way? — Hammamramma, Jun 12 '20 at 12:46

score 0 · Answer 2 · answered Jun 16 '20 at 18:16

thank you @pedrolarben for your solution, it helped me a lot but was incomplete!

what finally worked was the following

model_config = self.graph.get_config()
you need to change the inbound_nodes of the config of the layer not the inbound_nodes directly model_config['layers'][layer_index]['config']['layers'][index_of_dropout_layer]['inbound_nodes'][0][0][-1]['training'] = True
reload the model: (as stated in this answer The inverse of keras.models.Model.get_config() seems to be keras.models.Model.from_config(), not keras.models.model_from_config()):

model = Model.from_config(model_config)

And finally you need to load the weights again, otherwise your model is randomly initialized model.load_weights(self.graph_path)

Note: this works for the keras implementation of efficientNet.

ade sueb · Answer 3 · 2021-12-23T08:08:22.267

You can follow this link to create custom layer, also there is an example using custom dropout Layer, so that you can manipulate training values during training also inference time to make dropout running on inference time

class CustomDropout(keras.layers.Layer):
def __init__(self, rate, **kwargs):
    super(CustomDropout, self).__init__(**kwargs)
    self.rate = rate

def call(self, inputs, training=None):
    if training: #you can remove this line,, so that you can use dropout on inference time
        return tf.nn.dropout(inputs, rate=self.rate)
    return inputs

Is there a way to activate dropout during inference in Keras, while freezing batch norm layers

3 Answers3