[FIXED] How to log a tensorflow model with mlflow.tensorflow.log_model (error module 'tensorflow._api.v2.saved_model' has no attribute 'tag_constants')

Issue

I am trying to log a trained model with MLFlow using mlflow.tensorflow.log_model.

After training a simple sequential tf model

history = binary_model.fit(train_ds, validation_data=val_ds, epochs=num_epochs)

I am trying to log it:

    from tensorflow.python.saved_model import signature_constants
    tag=[tf.saved_model.tag_constants.SERVING]
    key=signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY

    mlflow.tensorflow.log_model(tf_saved_model_dir=saved_model_path,
                                tf_meta_graph_tags=tag,
                                tf_signature_def_key=key,
                                artifact_path="tf-models",
                                registered_model_name=model_name)

but I get the error:

    AttributeError                            Traceback (most recent call last)
    /var/folders/2k/g7p7j2gx6v54vkwv3v401h2m0000gn/T/ipykernel_73638/562549064.py in <module>
          1 from tensorflow.python.saved_model import signature_constants
    ----> 2 tag=[tf.saved_model.tag_constants.SERVING]
          3 key=signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY
          4 
          5 mlflow.tensorflow.log_model(tf_saved_model_dir=saved_model_path,

    AttributeError: module 'tensorflow._api.v2.saved_model' has no attribute 'tag_constants'

Any idea how to get the tags and keys correctly from the model to log it in MLFlow?

Many thanks in advance!

Solution

The tag_constants is in tf.compat.v1.saved_model.

To resolve the error replace this line

tag=[tf.saved_model.tag_constants.SERVING]

with this

tag=[tf.compat.v1.saved_model.tag_constants.SERVING]

Please refer this for more details.

Answered By - Tfer3