Search code examples
keraslstmkeras-layerdropout

How to use MC Dropout on a variational dropout LSTM layer on keras?

I'm currently trying to set up a (LSTM) recurrent neural network with Keras (tensorflow backend). I would like to use variational dropout with MC Dropout on it. I believe that variational dropout is already implemented with the option "recurrent_dropout" of the LSTM layer but I don't find any way to set a "training" flag to put on to true like a classical Dropout layer.

Solution

  • This is quite easy in Keras, first you need to define a function that takes both model input and the learning_phase:

    import keras.backend as K
f = K.function([model.layers[0].input, K.learning_phase()],
               [model.layers[-1].output])

    For a Functional API model with multiple inputs/outputs you can use:

    f = K.function([model.inputs, K.learning_phase()],
               [model.outputs])

    Then you can call this function like f([input, 1]) and this will tell Keras to enable the learning phase during this call, executing Dropout. Then you can call this function multiple times and combine the predictions to estimate uncertainty.