Dynamic switching of dropout in Keras/Tensorflow

Question

I am building a reinforcement learning algorithm in Tensorflow and I would like to be able to dynamically turn dropout off and then on within one single call to session.run().

Rationale: I need to (1) do a forward pass w/o dropout to calculate the targets; and (2) do a training step with the generated targets. If I execute these two steps in different calls to session.run(), everything is ok. But I would like to do it with one single call to session.run() (using tf.stop_gradients(targets)).

After trying several solutions w/o much success, I landed on a solution where I replace the learning_phase placeholder used by Keras with a variable (since placeholders are tensors and do not allow assignments) and use a custom layer to set that variable to True or False as desired. This solution is shown in the code below. Getting the value of either m1 or m2 separately (e.g., running sess.run(m1, feed_dict={ph:np.ones((1,1))})works as expected w/o error. However, getting the value of m3, or getting the values of m1 and m2 simultaneously, works sometimes and sometimes not (and the error message is uninformative).

Do you know what I am doing wrong or a better way to do what I want?

EDIT: The code shows a toy example. In reality I have a single model and I need to run two forward passes (one with dropout off and the other with dropout on) and one backward pass. And I want to do all this it w/o returning to python.

from tensorflow.keras.layers import Dropout, Dense, Input, Layer
from tensorflow.python.keras import backend as K
from tensorflow.keras import Model
import tensorflow as tf
import numpy as np

class DropoutSwitchLayer(Layer):
  def __init__(self, stateful=True, **kwargs):
    self.stateful = stateful
    self.supports_masking = True
    super(DropoutSwitchLayer, self).__init__(**kwargs)

  def build(self, input_shape):
    self.lph = tf.Variable(True, dtype=tf.bool, name="lph", trainable=False)
    K._GRAPH_LEARNING_PHASES[tf.get_default_graph()] = self.lph
    super(DropoutSwitchLayer, self).build(input_shape)

  def call(self, inputs, mask=None):
    data_input, training = inputs
    op = self.lph.assign(training[0], use_locking=True)
    # ugly trick here to make the layer work
    data_input = data_input + tf.multiply(tf.cast(op, dtype=tf.float32), 0.0)
    return data_input

  def compute_output_shape(self, input_shape):
    return input_shape[0]


dropout_on = np.array([True], dtype=np.bool)
dropout_off = np.array([False], dtype=np.bool)
input_ph = tf.placeholder(tf.float32, shape=(None, 1))

drop = Input(shape=(), dtype=tf.bool)
input = Input(shape=(1,))
h = DropoutSwitchLayer()([input, drop])
h = Dense(1)(h)
h = Dropout(0.5)(h)
o = Dense(1)(h)
m = Model(inputs=[input, drop], outputs=o)

m1 = m([input_ph, dropout_on])
m2 = m([input_ph, dropout_off])
m3 = m([m2, dropout_on])

sess = tf.Session()
K.set_session(sess)
sess.run(tf.global_variables_initializer())

EDIT 2: Daniel Möller's solution below works when using a Dropout layer, but what if using dropout inside an LSTM layer?

input = Input(shape=(1,))
h = Dense(1)(input)
h = RepeatVector(2)(h)
h = LSTM(1, dropout=0.5, recurrent_dropout=0.5)(h)
o = Dense(1)(h)

Answer 1

How about this :

class CustomDropout(tf.keras.layers.Layer):
    def __init__(self):
        super(CustomDropout, self).__init__()
        self.dropout1= Dropout(0.5)
        self.dropout2= Dropout(0.1)

    def call(self, inputs):
       if xxx:
           return self.dropout1(inputs)
       else:
           return self.dropout2(inputs)

Answer 2

It turns out Keras supports, out of the box, what I want to do. Using the training argument in the call to the Dropout/LSTM layer, in combination with Daniel Möller's approach to build the model (thanks!), does the trick.

In the code below (just a toy example), o1 and o3 should be equal and different than o2

from tensorflow.keras.layers import Dropout, Dense, Input, Lambda, Layer, Add, RepeatVector, LSTM
from tensorflow.python.keras import backend as K
from tensorflow.keras import Model
import tensorflow as tf
import numpy as np

repeat = RepeatVector(2)
lstm = LSTM(1, dropout=0.5, recurrent_dropout=0.5)

#Forward pass with dropout disabled
next_state = tf.placeholder(tf.float32, shape=(None, 1), name='next_state')
h = repeat(next_state)
# Use training to disable dropout
o1 = lstm(h, training=False)
target1 = tf.stop_gradient(o1)

#Forward pass with dropout enabled
state = tf.placeholder(tf.float32, shape=(None, 1), name='state')
h = repeat(state)
o2 = lstm(h, training=True)
target2 = tf.stop_gradient(o2)

#Forward pass with dropout disabled
ph3 = tf.placeholder(tf.float32, shape=(None, 1), name='ph3')
h = repeat(ph3)
o3 = lstm(h, training=False)

loss = target1 + target2 - o3
opt = tf.train.GradientDescentOptimizer(0.1)
train = opt.minimize(loss)

sess = tf.Session()
K.set_session(sess)
sess.run(tf.global_variables_initializer())

data = np.ones((1,1))
sess.run([o1, o2, o3], feed_dict={next_state:data, state:data, ph3:data})

Answer 3

Why not make a single continuous model?

#layers
inputs = Input(shape(1,))
dense1 = Dense(1)
dense2 = Dense(1)

#no drop pass:
h = dense1(inputs)
o = dense2(h)
#optionally:
o = Lambda(lambda x: K.stop_gradient(x))(o)

#drop pass:
h = dense1(o)
h = Dropout(.5)(h)
h = dense2(h)

modelOnlyFinalOutput = Model(inputs,h)
modelOnlyNonDrop = Model(inputs,o)
modelBothOutputs = Model(inputs, [o,h])

Choose one for training:

model.fit(x_train,y_train) #where y_train = [targets1, targets2] if using both outputs