tensorflow_serving/servables/tensorflow/testdata/saved_model_half_plus_two.py
def _generate_saved_model_for_half_plus_two(export_dir,
                                            tf2=False,
                                            as_text=False,
                                            as_tflite=False,
                                            as_tflite_with_sigdef=False,
                                            use_main_op=False,
                                            include_mlmd=False,
                                            device_type="cpu"):
  """Generates SavedModel for half plus two.

  Args:
    export_dir: The directory to which the SavedModel should be written.
    tf2: If True generates a SavedModel using native (non compat) TF2 APIs.
    as_text: Writes the SavedModel protocol buffer in text format to disk.
    as_tflite: Writes the Model in Tensorflow Lite format to disk.
    as_tflite_with_sigdef: Writes the Model with SignatureDefs in Tensorflow
      Lite format to disk.
    use_main_op: Whether to supply a main op during SavedModel build time.
    include_mlmd: Whether to include an MLMD key in the SavedModel.
    device_type: Device to force ops to run on.
  """
  if tf2:
    hp = HalfPlusTwoModel()
    tf.saved_model.save(hp, export_dir, signatures=hp.get_serving_signatures())
    return

  builder = tf.saved_model.builder.SavedModelBuilder(export_dir)

  device_name = "/cpu:0"
  if device_type == "gpu":
    device_name = "/gpu:0"

  with tf.Session(
      graph=tf.Graph(),
      config=tf.ConfigProto(log_device_placement=True)) as sess:
    with tf.device(device_name):
      # Set up the model parameters as variables to exercise variable loading
      # functionality upon restore.
      a = tf.Variable(0.5, name="a")
      b = tf.Variable(2.0, name="b")
      c = tf.Variable(3.0, name="c")

      # Create a placeholder for serialized tensorflow.Example messages to be
      # fed.
      serialized_tf_example = tf.placeholder(
          tf.string, name="tf_example", shape=[None])

      # parse_example only works on CPU
      with tf.device("/cpu:0"):
        tf_example = tf.parse_example(serialized_tf_example,
                                      _get_feature_spec())

      if as_tflite:
        # TFLite v1 converter does not support unknown shape.
        x = tf.ensure_shape(tf_example["x"], (1, 1), name="x")
      else:
        # Use tf.identity() to assign name
        x = tf.identity(tf_example["x"], name="x")

      if as_tflite_with_sigdef:
        # Resulting TFLite model will have input named "tflite_input".
        x = tf.ensure_shape(tf_example["x"], (1, 1), name="tflite_input")

      if device_type == "mkl":
        # Create a small convolution op to trigger MKL
        # The op will return 0s so this won't affect the
        # resulting calculation.
        o1 = tf.keras.layers.Conv2D(1, [1, 1])(tf.zeros((1, 16, 16, 1)))
        y = o1[0, 0, 0, 0] + tf.add(tf.multiply(a, x), b)
      else:
        y = tf.add(tf.multiply(a, x), b)

      y = tf.identity(y, name="y")

      if device_type == "mkl":
        # Create a small convolution op to trigger MKL
        # The op will return 0s so this won't affect the
        # resulting calculation.
        o2 = tf.keras.layers.Conv2D(1, [1, 1])(tf.zeros((1, 16, 16, 1)))
        y2 = o2[0, 0, 0, 0] + tf.add(tf.multiply(a, x), c)
      else:
        y2 = tf.add(tf.multiply(a, x), c)

      y2 = tf.identity(y2, name="y2")

      x2 = tf.identity(tf_example["x2"], name="x2")

      if device_type == "mkl":
        # Create a small convolution op to trigger MKL
        # The op will return 0s so this won't affect the
        # resulting calculation.
        o3 = tf.keras.layers.Conv2D(1, [1, 1])(tf.zeros((1, 16, 16, 1)))
        y3 = o3[0, 0, 0, 0] + tf.add(tf.multiply(a, x2), c)
      else:
        # Add separate constants for x2, to prevent optimizers like TF-TRT from
        # fusing the paths to compute y/y2 and y3 together.
        a2 = tf.Variable(0.5, name="a2")
        c2 = tf.Variable(3.0, name="c2")
        y3 = tf.add(tf.multiply(a2, x2), c2)

      y3 = tf.identity(y3, name="y3")

    assign_filename_op = _create_asset_file()

    predict_signature_def = _build_predict_signature(x, y)

    signature_def_map = {
        "regress_x_to_y":
            _build_regression_signature(serialized_tf_example, y),
        "regress_x_to_y2":
            _build_regression_signature(serialized_tf_example, y2),
        "regress_x2_to_y3":
            _build_regression_signature(x2, y3),
        "classify_x_to_y":
            _build_classification_signature(serialized_tf_example, y),
        tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
            _build_predict_signature(x, y)
    }
    # Initialize all variables and then save the SavedModel.
    sess.run(tf.global_variables_initializer())

    if as_tflite or as_tflite_with_sigdef:
      converter = tf.lite.TFLiteConverter.from_session(sess, [x], [y])
      tflite_model = converter.convert()
      if as_tflite_with_sigdef:
        k = tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY
        tflite_model = signature_def_utils.set_signature_defs(
            tflite_model, {k: predict_signature_def})
      open(export_dir + "/model.tflite", "wb").write(tflite_model)
    else:
      if use_main_op:
        builder.add_meta_graph_and_variables(
            sess, [tf.saved_model.tag_constants.SERVING],
            signature_def_map=signature_def_map,
            assets_collection=tf.get_collection(tf.GraphKeys.ASSET_FILEPATHS),
            main_op=tf.group(tf.saved_model.main_op.main_op(),
                             assign_filename_op))
      else:
        builder.add_meta_graph_and_variables(
            sess, [tf.saved_model.tag_constants.SERVING],
            signature_def_map=signature_def_map,
            assets_collection=tf.get_collection(tf.GraphKeys.ASSET_FILEPATHS),
            main_op=tf.group(assign_filename_op))

  if not as_tflite:
    builder.save(as_text)

  if include_mlmd:
    _write_mlmd(export_dir, "test_mlmd_uuid")
Generates SavedModel for half plus two.

  Args:
    export_dir: The directory to which the SavedModel should be written.
    tf2: If True generates a SavedModel using native (non compat) TF2 APIs.
    as_text: Writes the SavedModel protocol buffer in text format to disk.
    as_tflite: Writes the Model in Tensorflow Lite format to disk.
    as_tflite_with_sigdef: Writes the Model with SignatureDefs in Tensorflow
      Lite format to disk.
    use_main_op: Whether to supply a main op during SavedModel build time.
    include_mlmd: Whether to include an MLMD key in the SavedModel.
    device_type: Device to force ops to run on.
  
Generates SavedModel for half plus two.
Generates SavedModel for half plus two .
def _generate_saved_model_for_half_plus_two(export_dir,
                                            tf2=False,
                                            as_text=False,
                                            as_tflite=False,
                                            as_tflite_with_sigdef=False,
                                            use_main_op=False,
                                            include_mlmd=False,
                                            device_type="cpu"):
  if tf2:
    hp = HalfPlusTwoModel()
    tf.saved_model.save(hp, export_dir, signatures=hp.get_serving_signatures())
    return
  builder = tf.saved_model.builder.SavedModelBuilder(export_dir)
  device_name = "/cpu:0"
  if device_type == "gpu":
    device_name = "/gpu:0"
  with tf.Session(
      graph=tf.Graph(),
      config=tf.ConfigProto(log_device_placement=True)) as sess:
    with tf.device(device_name):
      a = tf.Variable(0.5, name="a")
      b = tf.Variable(2.0, name="b")
      c = tf.Variable(3.0, name="c")
      serialized_tf_example = tf.placeholder(
          tf.string, name="tf_example", shape=[None])
      with tf.device("/cpu:0"):
        tf_example = tf.parse_example(serialized_tf_example,
                                      _get_feature_spec())
      if as_tflite:
        x = tf.ensure_shape(tf_example["x"], (1, 1), name="x")
      else:
        x = tf.identity(tf_example["x"], name="x")
      if as_tflite_with_sigdef:
        x = tf.ensure_shape(tf_example["x"], (1, 1), name="tflite_input")
      if device_type == "mkl":
        o1 = tf.keras.layers.Conv2D(1, [1, 1])(tf.zeros((1, 16, 16, 1)))
        y = o1[0, 0, 0, 0] + tf.add(tf.multiply(a, x), b)
      else:
        y = tf.add(tf.multiply(a, x), b)
      y = tf.identity(y, name="y")
      if device_type == "mkl":
        o2 = tf.keras.layers.Conv2D(1, [1, 1])(tf.zeros((1, 16, 16, 1)))
        y2 = o2[0, 0, 0, 0] + tf.add(tf.multiply(a, x), c)
      else:
        y2 = tf.add(tf.multiply(a, x), c)
      y2 = tf.identity(y2, name="y2")
      x2 = tf.identity(tf_example["x2"], name="x2")
      if device_type == "mkl":
        o3 = tf.keras.layers.Conv2D(1, [1, 1])(tf.zeros((1, 16, 16, 1)))
        y3 = o3[0, 0, 0, 0] + tf.add(tf.multiply(a, x2), c)
      else:
        a2 = tf.Variable(0.5, name="a2")
        c2 = tf.Variable(3.0, name="c2")
        y3 = tf.add(tf.multiply(a2, x2), c2)
      y3 = tf.identity(y3, name="y3")
    assign_filename_op = _create_asset_file()
    predict_signature_def = _build_predict_signature(x, y)
    signature_def_map = {
        "regress_x_to_y":
            _build_regression_signature(serialized_tf_example, y),
        "regress_x_to_y2":
            _build_regression_signature(serialized_tf_example, y2),
        "regress_x2_to_y3":
            _build_regression_signature(x2, y3),
        "classify_x_to_y":
            _build_classification_signature(serialized_tf_example, y),
        tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
            _build_predict_signature(x, y)
    }
    sess.run(tf.global_variables_initializer())
    if as_tflite or as_tflite_with_sigdef:
      converter = tf.lite.TFLiteConverter.from_session(sess, [x], [y])
      tflite_model = converter.convert()
      if as_tflite_with_sigdef:
        k = tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY
        tflite_model = signature_def_utils.set_signature_defs(
            tflite_model, {k: predict_signature_def})
      open(export_dir + "/model.tflite", "wb").write(tflite_model)
    else:
      if use_main_op:
        builder.add_meta_graph_and_variables(
            sess, [tf.saved_model.tag_constants.SERVING],
            signature_def_map=signature_def_map,
            assets_collection=tf.get_collection(tf.GraphKeys.ASSET_FILEPATHS),
            main_op=tf.group(tf.saved_model.main_op.main_op(),
                             assign_filename_op))
      else:
        builder.add_meta_graph_and_variables(
            sess, [tf.saved_model.tag_constants.SERVING],
            signature_def_map=signature_def_map,
            assets_collection=tf.get_collection(tf.GraphKeys.ASSET_FILEPATHS),
            main_op=tf.group(assign_filename_op))
  if not as_tflite:
    builder.save(as_text)
  if include_mlmd:
    _write_mlmd(export_dir, "test_mlmd_uuid")
def _generate_saved_model_for_half_plus_two ( export_dir , tf2 = False , as_text = False , as_tflite = False , as_tflite_with_sigdef = False , use_main_op = False , include_mlmd = False , device_type = "cpu" ) : if tf2 : hp = HalfPlusTwoModel ( ) tf . saved_model . save ( hp , export_dir , signatures = hp . get_serving_signatures ( ) ) return builder = tf . saved_model . builder . SavedModelBuilder ( export_dir ) device_name = "/cpu:0" if device_type == "gpu" : device_name = "/gpu:0" with tf . Session ( graph = tf . Graph ( ) , config = tf . ConfigProto ( log_device_placement = True ) ) as sess : with tf . device ( device_name ) : a = tf . Variable ( 0.5 , name = "a" ) b = tf . Variable ( 2.0 , name = "b" ) c = tf . Variable ( 3.0 , name = "c" ) serialized_tf_example = tf . placeholder ( tf . string , name = "tf_example" , shape = [ None ] ) with tf . device ( "/cpu:0" ) : tf_example = tf . parse_example ( serialized_tf_example , _get_feature_spec ( ) ) if as_tflite : x = tf . ensure_shape ( tf_example [ "x" ] , ( 1 , 1 ) , name = "x" ) else : x = tf . identity ( tf_example [ "x" ] , name = "x" ) if as_tflite_with_sigdef : x = tf . ensure_shape ( tf_example [ "x" ] , ( 1 , 1 ) , name = "tflite_input" ) if device_type == "mkl" : o1 = tf . keras . layers . Conv2D ( 1 , [ 1 , 1 ] ) ( tf . zeros ( ( 1 , 16 , 16 , 1 ) ) ) y = o1 [ 0 , 0 , 0 , 0 ] + tf . add ( tf . multiply ( a , x ) , b ) else : y = tf . add ( tf . multiply ( a , x ) , b ) y = tf . identity ( y , name = "y" ) if device_type == "mkl" : o2 = tf . keras . layers . Conv2D ( 1 , [ 1 , 1 ] ) ( tf . zeros ( ( 1 , 16 , 16 , 1 ) ) ) y2 = o2 [ 0 , 0 , 0 , 0 ] + tf . add ( tf . multiply ( a , x ) , c ) else : y2 = tf . add ( tf . multiply ( a , x ) , c ) y2 = tf . identity ( y2 , name = "y2" ) x2 = tf . identity ( tf_example [ "x2" ] , name = "x2" ) if device_type == "mkl" : o3 = tf . keras . layers . Conv2D ( 1 , [ 1 , 1 ] ) ( tf . zeros ( ( 1 , 16 , 16 , 1 ) ) ) y3 = o3 [ 0 , 0 , 0 , 0 ] + tf . add ( tf . multiply ( a , x2 ) , c ) else : a2 = tf . Variable ( 0.5 , name = "a2" ) c2 = tf . Variable ( 3.0 , name = "c2" ) y3 = tf . add ( tf . multiply ( a2 , x2 ) , c2 ) y3 = tf . identity ( y3 , name = "y3" ) assign_filename_op = _create_asset_file ( ) predict_signature_def = _build_predict_signature ( x , y ) signature_def_map = { "regress_x_to_y" : _build_regression_signature ( serialized_tf_example , y ) , "regress_x_to_y2" : _build_regression_signature ( serialized_tf_example , y2 ) , "regress_x2_to_y3" : _build_regression_signature ( x2 , y3 ) , "classify_x_to_y" : _build_classification_signature ( serialized_tf_example , y ) , tf . saved_model . signature_constants . DEFAULT_SERVING_SIGNATURE_DEF_KEY : _build_predict_signature ( x , y ) } sess . run ( tf . global_variables_initializer ( ) ) if as_tflite or as_tflite_with_sigdef : converter = tf . lite . TFLiteConverter . from_session ( sess , [ x ] , [ y ] ) tflite_model = converter . convert ( ) if as_tflite_with_sigdef : k = tf . saved_model . signature_constants . DEFAULT_SERVING_SIGNATURE_DEF_KEY tflite_model = signature_def_utils . set_signature_defs ( tflite_model , { k : predict_signature_def } ) open ( export_dir + "/model.tflite" , "wb" ) . write ( tflite_model ) else : if use_main_op : builder . add_meta_graph_and_variables ( sess , [ tf . saved_model . tag_constants . SERVING ] , signature_def_map = signature_def_map , assets_collection = tf . get_collection ( tf . GraphKeys . ASSET_FILEPATHS ) , main_op = tf . group ( tf . saved_model . main_op . main_op ( ) , assign_filename_op ) ) else : builder . add_meta_graph_and_variables ( sess , [ tf . saved_model . tag_constants . SERVING ] , signature_def_map = signature_def_map , assets_collection = tf . get_collection ( tf . GraphKeys . ASSET_FILEPATHS ) , main_op = tf . group ( assign_filename_op ) ) if not as_tflite : builder . save ( as_text ) if include_mlmd : _write_mlmd ( export_dir , "test_mlmd_uuid" )
Generates SavedModel for half plus two.
Generates SavedModel for half plus two .
"""Generates SavedModel for half plus two.

  Args:
    export_dir: The directory to which the SavedModel should be written.
    tf2: If True generates a SavedModel using native (non compat) TF2 APIs.
    as_text: Writes the SavedModel protocol buffer in text format to disk.
    as_tflite: Writes the Model in Tensorflow Lite format to disk.
    as_tflite_with_sigdef: Writes the Model with SignatureDefs in Tensorflow
      Lite format to disk.
    use_main_op: Whether to supply a main op during SavedModel build time.
    include_mlmd: Whether to include an MLMD key in the SavedModel.
    device_type: Device to force ops to run on.
  """ # Set up the model parameters as variables to exercise variable loading # functionality upon restore. # Create a placeholder for serialized tensorflow.Example messages to be # fed. # parse_example only works on CPU # TFLite v1 converter does not support unknown shape. # Use tf.identity() to assign name # Resulting TFLite model will have input named "tflite_input". # Create a small convolution op to trigger MKL # The op will return 0s so this won't affect the # resulting calculation. # Create a small convolution op to trigger MKL # The op will return 0s so this won't affect the # resulting calculation. # Create a small convolution op to trigger MKL # The op will return 0s so this won't affect the # resulting calculation. # Add separate constants for x2, to prevent optimizers like TF-TRT from # fusing the paths to compute y/y2 and y3 together. # Initialize all variables and then save the SavedModel.
export_dir
The directory to which the SavedModel should be written.
The directory to which the SavedModel should be written .
tf2
If True generates a SavedModel using native (non compat) TF2 APIs.
If True generates a SavedModel using native ( non compat ) TF2 APIs .
as_text
Writes the SavedModel protocol buffer in text format to disk.
Writes the SavedModel protocol buffer in text format to disk .
as_tflite
Writes the Model in Tensorflow Lite format to disk.
Writes the Model in Tensorflow Lite format to disk .
as_tflite_with_sigdef
Writes the Model with SignatureDefs in Tensorflow
Lite format to disk.
Writes the Model with SignatureDefs in Tensorflow Lite format to disk .
use_main_op
Whether to supply a main op during SavedModel build time.
Whether to supply a main op during SavedModel build time .
include_mlmd
Whether to include an MLMD key in the SavedModel.
Whether to include an MLMD key in the SavedModel .
device_type
Device to force ops to run on.
Device to force ops to run on .