How to create Keras model with optional inputs

Question

I'm looking for a way to create a Keras model with optional inputs. In raw TensorFlow, you can create placeholders with optional inputs as follows:

import numpy as np
import tensorflow as tf


def main():
    required_input = tf.placeholder(
        tf.float32,
        shape=(None, 2),
        name='required_input')

    default_optional_input = tf.random_uniform(
        shape=(tf.shape(required_input)[0], 3))
    optional_input = tf.placeholder_with_default(
        default_optional_input,
        shape=(None, 3),
        name='optional_input')

    output = tf.concat((required_input, optional_input), axis=-1)

    with tf.Session() as session:
        with_optional_input_output_np = session.run(output, feed_dict={
            required_input: np.random.uniform(size=(4, 2)),
            optional_input: np.random.uniform(size=(4, 3)),
        })

        print(f"with optional input: {with_optional_input_output_np}")

        without_optional_input_output_np = session.run(output, feed_dict={
            required_input: np.random.uniform(size=(4, 2)),
        })

        print(f"without optional input: {without_optional_input_output_np}")


if __name__ == '__main__':
    main()

In a similar fashion, I would want to be able to have optional inputs for my Keras model. It seems like the tensor argument in the keras.layers.Input.__init__ might be what I'm looking for, but at least it doesn't work as I was expecting (i.e. the same way as tf.placeholder_with_default shown above). Here's an example that breaks:

import numpy as np
import tensorflow as tf
import tensorflow_probability as tfp


def create_model(output_size):
    required_input = tf.keras.layers.Input(
        shape=(13, ), dtype='float32', name='required_input')

    batch_size = tf.shape(required_input)[0]

    def sample_optional_input(inputs, batch_size=None):
        base_distribution = tfp.distributions.MultivariateNormalDiag(
            loc=tf.zeros(output_size),
            scale_diag=tf.ones(output_size),
            name='sample_optional_input')

        return base_distribution.sample(batch_size)

    default_optional_input = tf.keras.layers.Lambda(
        sample_optional_input,
        arguments={'batch_size': batch_size}
    )(None)

    optional_input = tf.keras.layers.Input(
        shape=(output_size, ),
        dtype='float32',
        name='optional_input',
        tensor=default_optional_input)

    concat = tf.keras.layers.Concatenate(axis=-1)(
        [required_input, optional_input])
    dense = tf.keras.layers.Dense(
        output_size, activation='relu')(concat)

    model = tf.keras.Model(
        inputs=[required_input, optional_input],
        outputs=[dense])

    return model


def main():
    model = create_model(output_size=3)

    required_input_np = np.random.normal(size=(4, 13))
    outputs_np = model.predict({'required_input': required_input_np})
    print(f"outputs_np: {outputs_np}")

    required_input = tf.random_normal(shape=(4, 13))
    outputs = model({'required_input': required_input})
    print(f"outputs: {outputs}")


if __name__ == '__main__':
    main()

The first call to the model.predict seems to give correct output, but for some reason, the direct call to model fails with the following error:

ValueError: Layer model expects 2 inputs, but it received 1 input tensors. Inputs received: []

Can the tensor argument in Input.__init__ be used to implement optional inputs for Keras model as in my example above? If yes, what should I change in my example to make it run correctly? If not, what is the expected way of creating optional inputs in Keras?

Answer 1

I really don't think it's possible without workarounds. Keras was not meant for that.

But, noticing that you are using two different session.run commands for each case, it seems that it should be easy to do it with two models. One model uses the optional input, the other doesn't. You choose which one to use the same way you choose which session.run() to call.

That said, you can use Input(tensor=...) or simply create the optional input inside a Lambda layer. Both things are fine. But don't use Input(shape=..., tensor=...), these are redundant arguments and sometimes Keras does not deal well with redundancies like this.

Ideally, keep all operations inside Lambda layers, even the tf.shape operation.

That said:

required_input = tf.keras.layers.Input(
    shape=(13, ), dtype='float32', name='required_input')

#needs the input for the case you want to pass it:
optional_input_when_used = tf.keras.layers.Input(shape=(output_size,))


#operations should be inside Lambda layers
batch_size = Lambda(lambda x: tf.shape(x)[0])(required_input)

#updated for using the batch size coming from lambda
#you didn't use "inputs" anywhere in this function
def sample_optional_input(batch_size):
    base_distribution = tfp.distributions.MultivariateNormalDiag(
        loc=tf.zeros(output_size),
        scale_diag=tf.ones(output_size),
        name='sample_optional_input')

    return base_distribution.sample(batch_size)

#updated for using the batch size as input
default_optional_input = tf.keras.layers.Lambda(sample_optional_input)(batch_size)

#let's skip the concat for now - notice I'm not "using" this layer yet
dense_layer = tf.keras.layers.Dense(output_size, activation='relu')

#you could create the rest of the model here if it's big, so you don't create it twice 
#(check the final section of this answer)

Model using passed input:

concat_when_used = tf.keras.layers.Concatenate(axis=-1)(
    [required_input, optional_input_when_used]
)

dense_when_used = dense_layer(concat_when_used)  
#or final_part_of_the_model(concat_when_used)     

model_when_used = Model([required_input, optional_input_when_used], dense_when_used)

Model not using the optional input:

concat_not_used = tf.keras.layers.Concatenate(axis=-1)(
    [required_input, default_optional_input]
)

dense_not_used = dense_layer(concat_not_used) 
#or final_part_of_the_model(concat_not_used)

model_not_used = Model(required_input, dense_not_used)

It's ok to create two models like this and choose one to use (both models share the final layers, so they will always be trained together)

Now, at the point you choose which session.run, now you will choose which model to use:

model_when_used.predict([x1, x2])
model_when_used.fit([x1,x2], y)

model_not_used.predict(x)
model_not_used.fit(x, y)

How to create a shared final part?

If your final part is big, you will not want to call everything twice to create two models. In this case, create a final model first:

input_for_final = Input(shape_after_concat)
out = Dense(....)(input_for_final)
out = Dense(....)(out)
out = Dense(....)(out)
.......
final_part_of_the_model = Model(input_for_final, out)

Then use this final part in previous answer.

dense_when_used = final_part_of_the_model(concat_when_used)
dense_not_used = final_part_of_the_model(concat_not_used)