MC dropout - fixing dropped weights at test time to sample smooth functions

[charliekirkwood]

It has been explained in issue 9412 how to keep dropout on at test time in order to sample from the approximate Bayesian posterior that MC dropout provides.

However, is it possible in Keras to fix which weights are dropped for any set of predictions (e.g. across the range of inputs) so that each call of predict(model, data) can provide predictions from a single smooth function, rather than the current setup of predicting from a separate function for each individual input?

Doing so would allow visualisations as in the following:

http://www.cs.ox.ac.uk/people/yarin.gal/website/blog_2248.html

As mentioned in this blog written by the inventor of MC dropout, fixing the dropped weights for all test inputs make better visualization.

Does anyone have a solution for fixing the dropout weights using the keras dropout?

Originally posted by @hjliag in #9412 (comment)

[mhlg]

@charliekirkwood while it is possible as you noted to keep dropout during inference by calling the layer with Layer(x, training=True) this in my opinion do not solve the initial issue.

There are circumstances where you do not want to use dropout during training or test and only use it during inference to model the prediction uncertainty.

My usage case is in LSTM

I thought to train a model where I set the traing=True and dropout=0
tf.keras.layers.LSTM( ..., dropout=0.0)(x, training=True)

and later at inference time update the dropout values
this will also enable you to repeat with different dropout rates.

inputs = # network inputs
model = #MyLSTMNet
predictions = []

for dropout_rate in [0.025, 0.05, 0.10]:
      for layer in model.layers:
               if hasattr(layer, 'dropout'):
                        setattr(layer,  'dropout',  dropout_rate)
  
      for i in range(100):
            predictions.append(model.predict(inputs))

but this results in throwing an error.
Can't set the attribute "dropout", likely because it conflicts with an existing read-only @Property of the object. Please choose a different name.

as per your request
Keras defines a function called reset_dropout_mask()
"""Reset the cached dropout masks if any.
This is important for the RNN layer to invoke this in it call() method so
that the cached mask is cleared before calling the cell.call(). The mask
should be cached across the timestep within the same batch, but shouldn't
be cached between batches. Otherwise it will introduce unreasonable bias
against certain index of data within the batch.
"""
that is invoked by class LSTM(recurrent.DropoutRNNCellMixin, recurrent.LSTM)
call() function so I think it can be done by subclassing the tf.keras.layers.LSTM and add some logic in the call() function to prevent the mask to be resetted when needed

[charliekirkwood]

@mhlg Thanks for your ideas, it sounds like a function to prevent the dropout mask from being reset for each prediction is what I'm after. From what I can see, no framework makes this easy (please correct me if I'm wrong - it seems like a good feature to have).

I also wondered if a custom dropout layer could be created based on tfp.layers.IndependentBernoulli, but I could not find how to adjust the rate for that layer, nor am I sure it can be frozen by calling (training = False).

[charliekirkwood]

For anyone else who may be wanting this functionality - I have managed to achieve it by using get_weights and set_weights - and manually applying a different random dropout mask (vector of samples from a binomial distribution with p = 1-dropout rate) to the weights.

Here is the result - which also helps explain what I was after in the first place - red is the true mean, points are the training data, and each blue line is a separate 'ensemble member' in the dropout neural network:

Thanks

[mhlg]

very nice can you give me some hint regarding the keras implementation ?

[charliekirkwood]

@mhlg yes, here is the relevant code section - I'm actually working in R but should be recognisable from Python too:

`n.sims <- 100

for(i in 1:n.sims){
bnnmodel %>% load_model_weights_hdf5("bnnweights.hdf5", by_name = FALSE, skip_mismatch = FALSE, reshape = FALSE)
weights <- bnnmodel %>% get_weights()

dropweights <- weights

maskl1 <- rbinom(dim(dropweights[[1]])[2], 1, 1-droprate) # generate dropout mask for kernel and bias of layer 1
maskl2 <- rbinom(dim(dropweights[[3]])[2], 1, 1-droprate) # generate dropout mask for kernel and bias of layer 2

dropweights[[1]] <- t(t(dropweights[[1]]) * maskl1)(1/(1-droprate))
dropweights[[2]] <- (dropweights[[2]] * maskl1)(1/(1-droprate))

dropweights[[3]] <- t(t(dropweights[[3]]) * maskl2)(1/(1-droprate))
dropweights[[4]] <- (dropweights[[4]] * maskl2)(1/(1-droprate))

bnnmodel %>% set_weights(dropweights)

preds <- as.data.table(data.frame(x = c(seq(-1.5,0, length.out = 100),surf$x,seq(1,2.5, length.out = 100)), draw = i))
preds[, y := as.numeric(tfd_mean(bnnmodel(as.matrix(x))))]

if(i == 1){
sim <- copy(preds)
}else{
sim <- rbindlist(list(sim, preds), use.names = FALSE, fill = FALSE, idcol = NULL)
}
}`

You'll want to save the trained model weights first (as "bnnweights.hdf5" in my case), and the layers you apply dropout to in dropweights will depend on the architecture of the neural network - the above is for a fully connected ANN with two hidden layers

[mhlg]

@charliekirkwood thanks!