Loading SavedModel generated by an estimator in TF v2

I trained a model with the estimator BoostedTreesClassifier from TF2 and saved it as a SavedModel. I'm having some issue to understand how to load this SavedModel in a different script and use it to compute some prediction. I found tons of code about how to do that with Keras or TF1, but none with TF2. To be clear, I'm trying to avoid using tf.compat.v1.xxx, as I imagine there is a better way to do that in TF2.

  • I was expecting that we could create an estimator from a SavedModel and query it directly, but could not find any method for that

  • I tried to go with loading the model and getting the graph:

    model = tf.saved_model.load(model_dir_path, ["serve"])
    graph = model.graph

But that looks really low level and I'm not sure how to execute the graph (beside using tf.compat.v1, which again, I'm trying to avoid)

Let me know if you have any idea, thanks so much!