" ValueError: Expecting KerasTensor which is from tf.keras.Input()". Error in prediction with dropout function

I am trying to predict uncertainty in a regression problem using Dropout during testing as per Yarin Gal's article. I created a class using Keras's backend function as provided by this stack overflow question's answer. The class takes a NN model as input and randomly drops neurons during testing to give a stochastic estimate rather than deterministic output for a time-series forecasting.

I create a simple encoder-decoder model as shown below for the forecasting with 0.1 dropout during training:

input_sequence = Input(shape=(lookback, train_x.shape[2]))
encoder = LSTM(128, return_sequences=False)(input_sequence)
r_vec = RepeatVector(forward_pred)(encoder)
decoder = LSTM(128, return_sequences=True, dropout=0.1)(r_vec) #maybe use dropout=0.1
output = TimeDistributed(Dense(train_y.shape[2], activation='linear'))(decoder)


# optimiser = optimizers.Adam(clipnorm=1)

enc_dec_model = Model(input_sequence, output)
enc_dec_model.compile(loss="mean_squared_error",
              optimizer="adam",
              metrics=['mean_squared_error'])
enc_dec_model.summary()

After that, I define and call the DropoutPrediction class.

# Define the class:

class KerasDropoutPrediction(object):
    def __init__(self ,model):
        self.f = K.function(
            [model.layers[0].input,
             K.learning_phase()],
            [model.layers[-1].output])
    def predict(self ,x, n_iter=10):
        result = []
        for _ in range(n_iter):
            result.append(self.f([x , 1]))
        result = np.array(result).reshape(n_iter ,x.shape[0] ,x.shape[1]).T
        return result

# Call the object:
kdp = KerasDropoutPrediction(enc_dec_model)
y_pred_do = kdp.predict(x_test,n_iter=100)
y_pred_do_mean = y_pred_do.mean(axis=1)

However, in the line kdp = KerasDropoutPrediction(enc_dec_model), when I call the LSTM model, I got the following error message which says the input has to be a Keras Tensor. Can anyone help me with this error?

Error Message:

ValueError: Found unexpected instance while processing input tensors for keras functional model. Expecting KerasTensor which is from tf.keras.Input() or output from keras layer call(). Got: 0

To activate Dropout at inference time, you simply have to specify training=True (TF>2.0) in the layer of interest (in the last LSTM layer in your case)

with training=False

inp = Input(shape=(10, 1))
x = LSTM(1, dropout=0.3)(inp, training=False)
m = Model(inp,x)
# m.compile(...)
# m.fit(...)

X = np.random.uniform(0,1, (1,10,1))

output = []
for i in range(0,100):
    output.append(m.predict(X)) # always the same

with training=True

inp = Input(shape=(10, 1))
x = LSTM(1, dropout=0.3)(inp, training=True)
m = Model(inp,x)
# m.compile(...)
# m.fit(...)

X = np.random.uniform(0,1, (1,10,1))

output = []
for i in range(0,100):
    output.append(m.predict(X)) # always different

In your example, this becomes:

input_sequence = Input(shape=(lookback, train_x.shape[2]))
encoder = LSTM(128, return_sequences=False)(input_sequence)
r_vec = RepeatVector(forward_pred)(encoder)
decoder = LSTM(128, return_sequences=True, dropout=0.1)(r_vec, training=True)
output = TimeDistributed(Dense(train_y.shape[2], activation='linear'))(decoder)

enc_dec_model = Model(input_sequence, output)
enc_dec_model.compile(
    loss="mean_squared_error",
    optimizer="adam",
    metrics=['mean_squared_error']
)

enc_dec_model.fit(train_x, train_y, epochs=10, batch_size=32)

and the KerasDropoutPrediction:

class KerasDropoutPrediction(object):
    def __init__(self, model):
        self.model = model
    def predict(self, X, n_iter=10):
        result = []
        for _ in range(n_iter):
            result.append(self.model.predict(X))
        result = np.array(result)
        return result

kdp = KerasDropoutPrediction(enc_dec_model)
y_pred_do = kdp.predict(test_x, n_iter=100)
y_pred_do_mean = y_pred_do.mean(axis=0)