def export_model(model, train_input_fn, eval_input_fn, tf_feat_cols, base_dir):
"""Export TensorFlow estimator (model).
Args:
model (tf.estimator.Estimator): Model to export.
train_input_fn (function): Training input function to create data receiver spec.
eval_input_fn (function): Evaluation input function to create data receiver spec.
tf_feat_cols (list(tf.feature_column)): Feature columns.
base_dir (str): Base directory to export the model.
Returns:
str: Exported model path
"""
tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
train_rcvr_fn = build_supervised_input_receiver_fn_from_input_fn(
train_input_fn)
eval_rcvr_fn = build_supervised_input_receiver_fn_from_input_fn(
eval_input_fn)
serve_rcvr_fn = tf.estimator.export.build_parsing_serving_input_receiver_fn(
tf.feature_column.make_parse_example_spec(tf_feat_cols))
rcvr_fn_map = {
tf.estimator.ModeKeys.TRAIN: train_rcvr_fn,
tf.estimator.ModeKeys.EVAL: eval_rcvr_fn,
tf.estimator.ModeKeys.PREDICT: serve_rcvr_fn,
}
exported_path = model.experimental_export_all_saved_models(
export_dir_base=base_dir, input_receiver_fn_map=rcvr_fn_map)
return exported_path.decode("utf-8")
def test_build_supervised_input_receiver_fn_from_input_fn(self):
def dummy_input_fn():
return ({"x": constant_op.constant([[1], [1]]),
"y": constant_op.constant(["hello", "goodbye"])},
constant_op.constant([[1], [1]]))
input_receiver_fn = export.build_supervised_input_receiver_fn_from_input_fn(
dummy_input_fn)
with ops.Graph().as_default():
input_receiver = input_receiver_fn()
self.assertEqual(set(["x", "y"]),
set(input_receiver.features.keys()))
self.assertIsInstance(input_receiver.labels, ops.Tensor)
self.assertEqual(set(["x", "y", "label"]),
set(input_receiver.receiver_tensors.keys()))
def test_build_supervised_input_receiver_fn_from_input_fn_args(self):
def dummy_input_fn(feature_key="x"):
return ({feature_key: constant_op.constant([[1], [1]]),
"y": constant_op.constant(["hello", "goodbye"])},
{"my_label": constant_op.constant([[1], [1]])})
input_receiver_fn = export.build_supervised_input_receiver_fn_from_input_fn(
dummy_input_fn, feature_key="z")
with ops.Graph().as_default():
input_receiver = input_receiver_fn()
self.assertEqual(set(["z", "y"]),
set(input_receiver.features.keys()))
self.assertEqual(set(["my_label"]),
set(input_receiver.labels.keys()))
self.assertEqual(set(["z", "y", "my_label"]),
set(input_receiver.receiver_tensors.keys()))
def dummy_supervised_receiver_fn(): return export.build_supervised_input_receiver_fn_from_input_fn(dummy_input_fn)