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...