Skip to content

Commit

Permalink
Explicitly import estimator from tensorflow as a separate import inst…
Browse files Browse the repository at this point in the history
…ead of

accessing it via tf.estimator and depend on the tensorflow estimator target.

PiperOrigin-RevId: 436547871
  • Loading branch information
tfx-copybara committed Mar 22, 2022
1 parent 06d1eb7 commit 3d193b1
Showing 1 changed file with 3 additions and 2 deletions.
5 changes: 3 additions & 2 deletions tfx_bsl/beam/run_inference_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,7 @@
from googleapiclient import discovery
from googleapiclient import http
import tensorflow as tf
from tensorflow.compat.v1 import estimator as tf_estimator
from tfx_bsl.beam import run_inference
from tfx_bsl.beam import test_helpers
from tfx_bsl.public.proto import model_spec_pb2
Expand Down Expand Up @@ -87,12 +88,12 @@ def _build_predict_model(self, model_path):
default_value=0)
}
serving_receiver = (
tf.compat.v1.estimator.export.build_parsing_serving_input_receiver_fn(
tf_estimator.export.build_parsing_serving_input_receiver_fn(
input_tensors)())
output_tensors = {'y': serving_receiver.features['x'] * 2}
sess = tf.compat.v1.Session()
sess.run(tf.compat.v1.initializers.global_variables())
signature_def = tf.compat.v1.estimator.export.PredictOutput(
signature_def = tf_estimator.export.PredictOutput(
output_tensors).as_signature_def(serving_receiver.receiver_tensors)
builder = tf.compat.v1.saved_model.builder.SavedModelBuilder(model_path)
builder.add_meta_graph_and_variables(
Expand Down

0 comments on commit 3d193b1

Please sign in to comment.