Skip to content

Commit e4a4085

Browse files
committed
add test for two tag sets
1 parent 484b835 commit e4a4085

File tree

1 file changed

+37
-0
lines changed

1 file changed

+37
-0
lines changed

coremltools/test/api/test_api_examples.py

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -292,6 +292,43 @@ def test_convert_from_saved_model_dir():
292292
mlmodel = ct.convert("./saved_model")
293293
mlmodel.save("./model.mlmodel")
294294

295+
@staticmethod
296+
def test_convert_from_two_tags_saved_model_dir(tmpdir):
297+
298+
import tensorflow as tf
299+
from tensorflow.compat.v1.saved_model import build_tensor_info
300+
from tensorflow.compat.v1.saved_model import signature_constants
301+
from tensorflow.compat.v1.saved_model import signature_def_utils
302+
303+
@tf.function
304+
def add(a, b):
305+
return a + b
306+
307+
c = add.get_concrete_function(tf.constant(21.0), tf.constant(21.0))
308+
309+
save_path = str(tmpdir)
310+
builder = tf.compat.v1.saved_model.Builder(save_path)
311+
312+
with tf.compat.v1.Session(graph=c.graph) as sess:
313+
tensor_info_a = build_tensor_info(c.graph.inputs[0])
314+
tensor_info_b = build_tensor_info(c.graph.inputs[1])
315+
tensor_info_y = build_tensor_info(c.graph.outputs[0])
316+
317+
prediction_signature = signature_def_utils.build_signature_def(
318+
inputs={'a': tensor_info_a, 'b': tensor_info_b},
319+
outputs={'output': tensor_info_y},
320+
method_name=signature_constants.PREDICT_METHOD_NAME)
321+
322+
builder.add_meta_graph_and_variables(sess, ["serve"],
323+
signature_def_map={
324+
signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
325+
prediction_signature,
326+
})
327+
328+
builder.add_meta_graph(["serve", "tpu"])
329+
builder.save()
330+
331+
ct.convert(save_path, source="tensorflow", tags=["serve"])
295332

296333
@staticmethod
297334
def test_keras_custom_layer_model():

0 commit comments

Comments
 (0)