save model issue in TF 2.0 (tf.saved_model.save)

[shawei3000]

I am new to TF 2.0, I tried to save model by " tf.saved_model.save(squad_m......", but always get errors, such as: " start_positions = inputs["start_positions"] KeyError: 'start_positions'". I am guessing this is because the use of subclassing of keras_model: "class ALBertQAModel(tf.keras.Model):" , could you confirm or help me understand if otherwise?
Thanks,
Jim

[shawei3000]

I tried adding following in the Model sub-class:
@tf.function(input_signature=[{"unique_ids": tf.TensorSpec(shape=[ None], dtype=tf.int32),
"input_ids": tf.TensorSpec(shape=[ None, max_seq], dtype=tf.int32),
"input_mask": tf.TensorSpec(shape=[ None, max_seq], dtype=tf.int32),
"segment_ids": tf.TensorSpec(shape=[ None, max_seq], dtype=tf.int32),
"cls_index": tf.TensorSpec(shape=[ None], dtype=tf.int32),
"p_mask": tf.TensorSpec(shape=[ None, max_seq], dtype=tf.float32)
}#, bool
])
But still have issue/error export model during prediction....

[kamalkraj]

@shawei3000
Post a simple code snippet , where i could reproduce the above error
Try saving using https://www.tensorflow.org/api_docs/python/tf/keras/models/save_model

[shawei3000]

Thanks, Yes, in order to reproduce the error:
a). add following in class ALBertQAModel(tf.keras.Model), right above "def call(self, inputs):"

@tf.function(input_signature=[{"unique_ids": tf.TensorSpec(shape=[ None], dtype=tf.int32),
"input_ids": tf.TensorSpec(shape=[ None, your_max_seq], dtype=tf.int32),
"input_mask": tf.TensorSpec(shape=[ None, your_max_seq], dtype=tf.int32),
"segment_ids": tf.TensorSpec(shape=[ None, your_max_seq], dtype=tf.int32),
"cls_index": tf.TensorSpec(shape=[ None], dtype=tf.int32),
"p_mask": tf.TensorSpec(shape=[ None, your_max_seq], dtype=tf.float32)
}#, bool
])

b). in function "predict_squad_customized", right before prediction, add:
tf.keras.experimental.export_saved_model(
squad_model, your-prefered_dir,
serving_only=True,
input_signature=[{"unique_ids": tf.TensorSpec(shape=[ None], dtype=tf.int32),
"input_ids": tf.TensorSpec(shape=[ None, your_max_seq], dtype=tf.int32),
"input_mask": tf.TensorSpec(shape=[ None, your_max_seq], dtype=tf.int32),
"segment_ids": tf.TensorSpec(shape=[ None, your_max_seq], dtype=tf.int32),
"cls_index": tf.TensorSpec(shape=[ None], dtype=tf.int32),
"p_mask": tf.TensorSpec(shape=[ None, your_max_seq], dtype=tf.float32)}#,
],
# custom_objects={'loss':get_loss_fn_v2}
)
c). run "run_squad.py" prediction only (mode=predict)

I am guessing 2 issues:
1. model (squad_model) input has a "training=False" element, which is not a tensor, thus difficult to add in signature..
2. even we asked only prediction graph, but seems the model saving still will be looking for loss function...