@@ -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