how to apply MC dropout to an LSTM network keras

Question

I have a simple LSTM network developped using keras:

model = Sequential()
model.add(LSTM(rnn_size,input_shape=(2,w),dropout = 0.25 , recurrent_dropout=0.25))
model.add(Dense(2))

I would like to apply the MC dropout method. How can I enable dropout in the test phase in order to compute the uncertainty?

Thanks.

Answer 1

Monte-Carlo Dropout is the use of dropout at inference time in order to add stochasticity to a network that can be used to generate a cohort of predictors/predictions that you can perform statistical analysis on. This is commonly used for bootstrapping confidence intervals.

Where you perform dropout in your sequential model is therefore important. If you add dropout to:

• an input layer - you're simulating how the model responds to uncertainty about the input data
• the LSTM layer itself (recurrently) - you're simulating how varying the weights of the LSTM itself can impact your predictions
• the final dense layer - you're simulating modifying the projection of the LSTMs predictions into the original problem space

So Marat's answer is correct, but it only tells you how to add an additional dropout layer, which won't work if you want the dropout applied recursively within the LSTM.

As you've identified, you can't just configure the LSTM layer to use dropout because it won't be applied at inference, so instead we can subclass the built-in LSTM layer and force it to always behave in training mode:

class MonteCarloLSTM(tf.keras.layers.LSTM):
   def call(self, inputs):
      return super().call(inputs, training=True)

We then use our new MonteCarloLSTM layer in place of LSTM when defining our Sequential model:

model.add(MonteCarloLSTM(..., recurrent_dropout=0.2))

Note that this will not only give you Monte-Carlo dropout at inference time but also use dropout regularisation at training time. If you want to avoid the latter then you need to:

1. Train your model with a vanilla LSTM layer
2. Save the weights of all layers in your Sequential model
3. Rebuild a new Sequential model with exactly the same structure except replacing the LSTM layer with your custom MonteCarloLSTM layer
4. Load the trained weights into your model
5. Run inference

Answer 2

Well, in order to enable dropout during test phase you can do something like this:

keras.layers.Dropout(0.5)(x, training=True)

Then you'll probably want to run it multiple times. If you don't care about the inference time just run forward pass multiple times and at the end calculate mean and variance of your output.