CMCDragonkai
CMCDragonkai

Reputation: 6372

Keras model with 2 inputs during training, but only 1 input during inference

I have a Keras model that requires 2 inputs during training (and has 2 outputs corresponding to 2 loss functions).

The 2 inputs and 2 outputs are pairwise connected to each other.

Therefore in inference I don't actually need to pass the second input, nor do I need the second output.

Is there a way to make Keras/tf.keras with the predict method accept only the first input and produce the first output and ignore the second input and second output.

I can create a zeroed numpy array for the second input, but I'm wondering if it's possible to reduce memory usage or computation.

Tensorflow should be able to do this since its graph is lazy. But is Keras capable of doing this?

Example:

# assume second_batch is not needed
second_batch = np.zeros(shape=first_batch.shape)
results = model.predict((first_batch, second_batch))
# i only care about results[0]
# not results[1]

Upvotes: 2

Views: 1537

Answers (1)

Srihari Humbarwadi
Srihari Humbarwadi

Reputation: 2642

You can always construct a new keras model using the same shared weights, and specifying the required input and output tensors.

import tensorflow as tf

print('TensorFlow:', tf.__version__)


input_a = tf.keras.Input(shape=[224, 224, 3], name='input_a')
input_b = tf.keras.Input(shape=[224, 224, 3], name='input_b')

resnet_model = tf.keras.applications.ResNet50(include_top=False, pooling='avg')

xa = resnet_model(input_a)
xb = resnet_model(input_b)

output_a = tf.keras.layers.Dense(10, name='output_a', activation='softmax')(xa)
output_b = tf.keras.layers.Dense(10, name='output_b', activation='softmax')(xb)

training_model = tf.keras.Model(inputs=[input_a, input_b], outputs=[output_a, output_b])

[print('Training Model Input:', x.name, x.shape) for x in training_model.inputs]
print('')
[print('Training Model Output:', x.name, x.shape) for x in training_model.outputs]
print('')


inference_model = tf.keras.Model(inputs=[input_a], outputs=[output_a])

[print('Inference Model Input:', x.name, x.shape) for x in inference_model.inputs]
[print('Inference Model Output:', x.name, x.shape) for x in inference_model.outputs]

image = tf.random.uniform([1, 224, 224, 3])
predictions = inference_model(image, training=False)
print('')
print('Predictions:', predictions)

Output:

TensorFlow: 2.3.0-dev20200625
Training Model Input: input_a:0 (None, 224, 224, 3)
Training Model Input: input_b:0 (None, 224, 224, 3)

Training Model Output: output_a/Softmax:0 (None, 10)
Training Model Output: output_b/Softmax:0 (None, 10)

Inference Model Input: input_a:0 (None, 224, 224, 3)
Inference Model Output: output_a/Softmax:0 (None, 10)

Predictions: tf.Tensor(
[[0.01937425 0.17703871 0.08633    0.06593429 0.18057525 0.03161139
  0.01154568 0.09730788 0.01927926 0.31100336]], shape=(1, 10), dtype=float32)

Upvotes: 3

Related Questions