Python tensorflow.python.saved_model.signature_constants.PREDICT_METHOD_NAME Examples

def build_signature(inputs, outputs):
  """Build the signature.

  Not using predic_signature_def in saved_model because it is replacing the
  tensor name, b/35900497.

  Args:
    inputs: a dictionary of tensor name to tensor
    outputs: a dictionary of tensor name to tensor
  Returns:
    The signature, a SignatureDef proto.
  """
  signature_inputs = {key: saved_model_utils.build_tensor_info(tensor)
                      for key, tensor in inputs.items()}
  signature_outputs = {key: saved_model_utils.build_tensor_info(tensor)
                       for key, tensor in outputs.items()}

  signature_def = signature_def_utils.build_signature_def(
      signature_inputs, signature_outputs,
      signature_constants.PREDICT_METHOD_NAME)

  return signature_def 

def build_signature(inputs, outputs):
  """Build the signature.

  Not using predic_signature_def in saved_model because it is replacing the
  tensor name, b/35900497.

  Args:
    inputs: a dictionary of tensor name to tensor
    outputs: a dictionary of tensor name to tensor
  Returns:
    The signature, a SignatureDef proto.
  """
  signature_inputs = {key: saved_model_utils.build_tensor_info(tensor)
                      for key, tensor in inputs.items()}
  signature_outputs = {key: saved_model_utils.build_tensor_info(tensor)
                       for key, tensor in outputs.items()}

  signature_def = signature_def_utils.build_signature_def(
      signature_inputs, signature_outputs,
      signature_constants.PREDICT_METHOD_NAME)

  return signature_def 

def _convert_named_signatures_to_signature_def(signatures):
  """Convert named signatures to object of type SignatureDef.

  Args:
    signatures: object of type manifest_pb2.Signatures()

  Returns:
    object of type SignatureDef which contains a converted version of named
    signatures from input signatures object

  Raises:
    RuntimeError: if input and output named signatures are not of type
    GenericSignature
  """
  signature_def = meta_graph_pb2.SignatureDef()
  input_signature = signatures.named_signatures[
      signature_constants.PREDICT_INPUTS]
  output_signature = signatures.named_signatures[
      signature_constants.PREDICT_OUTPUTS]
  # TODO(pdudnik): what if there are other signatures? Mimic cr/140900781 once
  # it is submitted.
  if (input_signature.WhichOneof("type") != "generic_signature" or
      output_signature.WhichOneof("type") != "generic_signature"):
    raise RuntimeError("Named input and output signatures can only be "
                       "up-converted if they are generic signature. "
                       "Input signature type is %s, output signature type is "
                       "%s" % (input_signature.WhichOneof("type"),
                               output_signature.WhichOneof("type")))

  signature_def.method_name = signature_constants.PREDICT_METHOD_NAME
  for key, val in input_signature.generic_signature.map.items():
    _add_input_to_signature_def(val.tensor_name, key, signature_def)
  for key, val in output_signature.generic_signature.map.items():
    _add_output_to_signature_def(val.tensor_name, key, signature_def)
  return signature_def 

def export_saved_model(sess, export_path, input_tensor, output_tensor):
    from tensorflow.python.saved_model import builder as saved_model_builder
    from tensorflow.python.saved_model import signature_constants
    from tensorflow.python.saved_model import signature_def_utils
    from tensorflow.python.saved_model import tag_constants
    from tensorflow.python.saved_model import utils
    builder = saved_model_builder.SavedModelBuilder(export_path)

    prediction_signature = signature_def_utils.build_signature_def(
        inputs={'images': utils.build_tensor_info(input_tensor)},
        outputs={
            'scores': utils.build_tensor_info(output_tensor)
        },
        method_name=signature_constants.PREDICT_METHOD_NAME)

    legacy_init_op = tf.group(
        tf.tables_initializer(), name='legacy_init_op')
    builder.add_meta_graph_and_variables(
        sess, [tag_constants.SERVING],
        signature_def_map={
            'predict_images':
                prediction_signature,
        },
        legacy_init_op=legacy_init_op)

    builder.save() 

def save_signature(self, directory):

        signature = signature_def_utils.build_signature_def(
            inputs={
                'input':
                saved_model_utils.build_tensor_info(self.input),
                'dropout_rate':
                saved_model_utils.build_tensor_info(self.dropout_rate)
            },
            outputs={
                'output': saved_model_utils.build_tensor_info(self.output)
            },
            method_name=signature_constants.PREDICT_METHOD_NAME)
        signature_map = {
            signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature
        }
        model_builder = saved_model_builder.SavedModelBuilder(directory)
        model_builder.add_meta_graph_and_variables(
            self.sess,
            tags=[tag_constants.SERVING],
            signature_def_map=signature_map,
            clear_devices=True)
        model_builder.save(as_text=False) 

def export_model(self, model_dir, global_step_val, last_checkpoint):
    """Exports the model so that it can used for batch predictions."""

    with self.graph.as_default():
      with tf.Session() as session:
        session.run(tf.global_variables_initializer())
        self.saver.restore(session, last_checkpoint)

        signature = signature_def_utils.build_signature_def(
            inputs=self.inputs,
            outputs=self.outputs,
            method_name=signature_constants.PREDICT_METHOD_NAME)

        signature_map = {signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
                         signature}

        model_builder = saved_model_builder.SavedModelBuilder(model_dir)
        model_builder.add_meta_graph_and_variables(session,
            tags=[tag_constants.SERVING],
            signature_def_map=signature_map,
            clear_devices=True)
        model_builder.save() 

def export_model():
    """Exports the model"""

    trained_checkpoint_prefix = 'linear_regression'

    loaded_graph = tf.Graph()
    with tf.Session(graph=loaded_graph) as sess:
        sess.run(tf.global_variables_initializer())

        # Restore from checkpoint:
        loader = tf.train.import_meta_graph(trained_checkpoint_prefix + '.meta')
        loader.restore(sess, trained_checkpoint_prefix)

        # Add signature:
        graph = tf.get_default_graph()
        inputs = tf.saved_model.utils.build_tensor_info(graph.get_tensor_by_name('X:0'))
        outputs = tf.saved_model.utils.build_tensor_info(graph.get_tensor_by_name('y_model:0'))

        signature = signature_def_utils.build_signature_def(inputs={'X': inputs},
                                                            outputs={'y_model': outputs},
                                                            method_name=signature_constants.PREDICT_METHOD_NAME)

        signature_map = {signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature}

        # Export model:
        builder = tf.saved_model.builder.SavedModelBuilder('./my_model')
        builder.add_meta_graph_and_variables(sess, signature_def_map=signature_map,
                                             tags=[tf.saved_model.tag_constants.SERVING])
        builder.save()


# Export the model: 

def export_model():
    """Exports the model"""

    trained_checkpoint_prefix = 'softmax_regression_model_mnist'

    loaded_graph = tf.Graph()
    with tf.Session(graph=loaded_graph) as sess:
        sess.run(tf.global_variables_initializer())

        # Restore from checkpoint
        loader = tf.train.import_meta_graph(trained_checkpoint_prefix + '.meta')
        loader.restore(sess, trained_checkpoint_prefix)

        # Add signature:
        graph = tf.get_default_graph()
        inputs = tf.saved_model.utils.build_tensor_info(graph.get_tensor_by_name('myInput:0'))
        outputs = tf.saved_model.utils.build_tensor_info(graph.get_tensor_by_name('myOutput:0'))

        signature = signature_def_utils.build_signature_def(inputs={'myInput': inputs},
                                                            outputs={'myOutput': outputs},
                                                            method_name=signature_constants.PREDICT_METHOD_NAME)

        signature_map = {signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature}

        # Export model:
        builder = tf.saved_model.builder.SavedModelBuilder('./my_model')
        builder.add_meta_graph_and_variables(sess, signature_def_map=signature_map,
                                             tags=[tf.saved_model.tag_constants.SERVING])
        builder.save()


# Export the model: 

def export_model(self, model_dir, global_step_val, last_checkpoint):
    """Exports the model so that it can used for batch predictions."""

    with self.graph.as_default():
      with tf.Session() as session:
        session.run(tf.global_variables_initializer())
        self.saver.restore(session, last_checkpoint)

        signature = signature_def_utils.build_signature_def(
            inputs=self.inputs,
            outputs=self.outputs,
            method_name=signature_constants.PREDICT_METHOD_NAME)

        signature_map = {
            signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: signature
        }

        model_builder = saved_model_builder.SavedModelBuilder(model_dir)
        model_builder.add_meta_graph_and_variables(
            session,
            tags=[tag_constants.SERVING],
            signature_def_map=signature_map,
            clear_devices=True)
        model_builder.save() 

def export_model(self, model_dir, global_step_val, last_checkpoint):
        """Exports the model so that it can used for batch predictions."""
        with self.graph.as_default():
            with tf.Session() as session:
                session.run(tf.global_variables_initializer())
                self.saver.restore(session, last_checkpoint)

                signature = signature_def_utils.build_signature_def(
                    inputs=self.inputs,
                    outputs=self.outputs,
                    method_name=signature_constants.PREDICT_METHOD_NAME)

                signature_map = {signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
                                     signature}

                model_builder = saved_model_builder.SavedModelBuilder(model_dir)
                model_builder.add_meta_graph_and_variables(session,
                                                           tags=[tag_constants.SERVING],
                                                           signature_def_map=signature_map,
                                                           clear_devices=True)
                model_builder.save() 

def export_model(self, model_dir, global_step_val, last_checkpoint):
    """Exports the model so that it can used for batch predictions."""

    with self.graph.as_default():
      with tf.Session(config=self.config) as session:
        session.run(tf.global_variables_initializer())
        self.saver.restore(session, last_checkpoint)

        signature = signature_def_utils.build_signature_def(
            inputs=self.inputs,
            outputs=self.outputs,
            method_name=signature_constants.PREDICT_METHOD_NAME)

        signature_map = {signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
                         signature}

        model_builder = saved_model_builder.SavedModelBuilder(model_dir)
        model_builder.add_meta_graph_and_variables(session,
            tags=[tag_constants.SERVING],
            signature_def_map=signature_map,
            clear_devices=True)
        model_builder.save() 

def export_model(self, model_dir, global_step_val, last_checkpoint):
    """Exports the model so that it can used for batch predictions."""

    with self.graph.as_default():
      with tf.Session(config=self.config) as session:
        session.run(tf.global_variables_initializer())
        self.saver.restore(session, last_checkpoint)

        signature = signature_def_utils.build_signature_def(
            inputs=self.inputs,
            outputs=self.outputs,
            method_name=signature_constants.PREDICT_METHOD_NAME)

        signature_map = {signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
                         signature}

        model_builder = saved_model_builder.SavedModelBuilder(model_dir)
        model_builder.add_meta_graph_and_variables(session,
            tags=[tag_constants.SERVING],
            signature_def_map=signature_map,
            clear_devices=True)
        model_builder.save() 

def export_model(self, model_dir, global_step_val, last_checkpoint):
    """Exports the model so that it can used for batch predictions."""

    with self.graph.as_default():
      with tf.Session() as session:
        session.run(tf.global_variables_initializer())
        self.saver.restore(session, last_checkpoint)

        signature = signature_def_utils.build_signature_def(
            inputs=self.inputs,
            outputs=self.outputs,
            method_name=signature_constants.PREDICT_METHOD_NAME)

        signature_map = {signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
                         signature}

        model_builder = saved_model_builder.SavedModelBuilder(model_dir)
        model_builder.add_meta_graph_and_variables(session,
            tags=[tag_constants.SERVING],
            signature_def_map=signature_map,
            clear_devices=True)
        model_builder.save() 

def predict_signature_def(inputs, outputs):
  """Creates prediction signature from given inputs and outputs.

  Args:
    inputs: dict of string to `Tensor`.
    outputs: dict of string to `Tensor`.

  Returns:
    A prediction-flavored signature_def.

  Raises:
    ValueError: If inputs or outputs is `None`.
  """
  if inputs is None or not inputs:
    raise ValueError('Prediction inputs cannot be None or empty.')
  if outputs is None or not outputs:
    raise ValueError('Prediction outputs cannot be None or empty.')

  signature_inputs = {key: utils.build_tensor_info(tensor)
                      for key, tensor in inputs.items()}
  signature_outputs = {key: utils.build_tensor_info(tensor)
                       for key, tensor in outputs.items()}

  signature_def = build_signature_def(
      signature_inputs, signature_outputs,
      signature_constants.PREDICT_METHOD_NAME)

  return signature_def 

def _is_valid_predict_signature(signature_def):
  """Determine whether the argument is a servable 'predict' SignatureDef."""
  if signature_def.method_name != signature_constants.PREDICT_METHOD_NAME:
    return False
  if not signature_def.inputs.keys():
    return False
  if not signature_def.outputs.keys():
    return False
  return True 

def export(model_version, model_dir, sess, inputs, y_op):
    """导出tensorflow_serving可用的模型（Saved Model方式）（推荐）
    prediction_signature必备的三个参数分别是输入inputs、输出outputs和方法名method_name，如果缺失方法名将会报错：“grpc.framework.interfaces.face.face.AbortionError: AbortionError(code=StatusCode.INTERNAL, details="Expected prediction signature method_name to be one of {tensorflow/serving/predict, tensorflow/serving/classify, tensorflow/serving/regress}. Was: ")”。每一个SavedModel关联着一个独立的checkpoint。每一个图元都绑定一个或多个标签，这些标签用来明确图元被加载的方式。标签只接受两种类型：serve或者train，保存时可以同时包含两个标签。其中tag_constants.SERVING = "serve"，tag_constants.TRAINING = "train"。模型用于TensorFlow Serving时，标签必须包含serve类型。如果标签只包含train类型，TensorFlow Serving加载模型时会报错：“E tensorflow_serving/core/aspired_versions_manager.cc:351] Servable {name: default version: 2} cannot be loaded: Not found: Could not find meta graph def matching supplied tags.”。定义signature_def_map时注意定义默认服务签名键，如果缺少则会报错：“grpc.framework.interfaces.face.face.AbortionError: AbortionError(code=StatusCode.FAILED_PRECONDITION, details="Default serving signature key not found.")”。
    """
    if model_version <= 0:
        print('Please specify a positive value for version number.')
        sys.exit()

    path = os.path.dirname(os.path.abspath(model_dir))
    if os.path.isdir(path) == False:
        logging.warning('Path (%s) not exists, making directories...', path)
        os.makedirs(path)

    export_path = os.path.join(
        compat.as_bytes(model_dir),
        compat.as_bytes(str(model_version)))

    if os.path.isdir(export_path) == True:
        logging.warning('Path (%s) exists, removing directories...', export_path)
        shutil.rmtree(export_path)

    builder = saved_model_builder.SavedModelBuilder(export_path)
    tensor_info_x = utils.build_tensor_info(inputs)
    tensor_info_y = utils.build_tensor_info(y_op)

    prediction_signature = signature_def_utils.build_signature_def(
        inputs={'x': tensor_info_x},
        outputs={'y': tensor_info_y},
        method_name=signature_constants.PREDICT_METHOD_NAME)

    builder.add_meta_graph_and_variables(
        sess,
        [tag_constants.SERVING],
        signature_def_map={
            'predict_text': prediction_signature,
            signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: prediction_signature
        })

    builder.save() 

def testPredictionSignatureDef(self):
    input1 = constant_op.constant("a", name="input-1")
    input2 = constant_op.constant("b", name="input-2")
    output1 = constant_op.constant("c", name="output-1")
    output2 = constant_op.constant("d", name="output-2")
    signature_def = signature_def_utils.predict_signature_def({
        "input-1": input1,
        "input-2": input2
    }, {"output-1": output1,
        "output-2": output2})

    self.assertEqual(signature_constants.PREDICT_METHOD_NAME,
                     signature_def.method_name)

    # Check inputs in signature def.
    self.assertEqual(2, len(signature_def.inputs))
    input1_tensor_info_actual = (signature_def.inputs["input-1"])
    self.assertEqual("input-1:0", input1_tensor_info_actual.name)
    self.assertEqual(types_pb2.DT_STRING, input1_tensor_info_actual.dtype)
    self.assertEqual(0, len(input1_tensor_info_actual.tensor_shape.dim))
    input2_tensor_info_actual = (signature_def.inputs["input-2"])
    self.assertEqual("input-2:0", input2_tensor_info_actual.name)
    self.assertEqual(types_pb2.DT_STRING, input2_tensor_info_actual.dtype)
    self.assertEqual(0, len(input2_tensor_info_actual.tensor_shape.dim))

    # Check outputs in signature def.
    self.assertEqual(2, len(signature_def.outputs))
    output1_tensor_info_actual = (signature_def.outputs["output-1"])
    self.assertEqual("output-1:0", output1_tensor_info_actual.name)
    self.assertEqual(types_pb2.DT_STRING, output1_tensor_info_actual.dtype)
    self.assertEqual(0, len(output1_tensor_info_actual.tensor_shape.dim))
    output2_tensor_info_actual = (signature_def.outputs["output-2"])
    self.assertEqual("output-2:0", output2_tensor_info_actual.name)
    self.assertEqual(types_pb2.DT_STRING, output2_tensor_info_actual.dtype)
    self.assertEqual(0, len(output2_tensor_info_actual.tensor_shape.dim)) 

def _convert_named_signatures_to_signature_def(signatures):
  """Convert named signatures to object of type SignatureDef.

  Args:
    signatures: object of type manifest_pb2.Signatures()

  Returns:
    object of type SignatureDef which contains a converted version of named
    signatures from input signatures object

  Raises:
    RuntimeError: if input and output named signatures are not of type
    GenericSignature
  """
  signature_def = meta_graph_pb2.SignatureDef()
  input_signature = signatures.named_signatures[
      signature_constants.PREDICT_INPUTS]
  output_signature = signatures.named_signatures[
      signature_constants.PREDICT_OUTPUTS]
  # TODO(pdudnik): what if there are other signatures? Mimic cr/140900781 once
  # it is submitted.
  if (input_signature.WhichOneof("type") != "generic_signature" or
      output_signature.WhichOneof("type") != "generic_signature"):
    raise RuntimeError("Named input and output signatures can only be "
                       "up-converted if they are generic signature. "
                       "Input signature type is %s, output signature type is "
                       "%s" % (input_signature.WhichOneof("type"),
                               output_signature.WhichOneof("type")))

  signature_def.method_name = signature_constants.PREDICT_METHOD_NAME
  for key, val in input_signature.generic_signature.map.items():
    _add_input_to_signature_def(val.tensor_name, key, signature_def)
  for key, val in output_signature.generic_signature.map.items():
    _add_output_to_signature_def(val.tensor_name, key, signature_def)
  return signature_def 

def export_model(self, model_dir, global_step_val, last_checkpoint):
    """Exports the model so that it can used for batch predictions."""

    with self.graph.as_default():
      with tf.Session() as session:
        session.run(tf.global_variables_initializer())
        self.saver.restore(session, last_checkpoint)

        signature = signature_def_utils.build_signature_def(
            inputs=self.inputs,
            outputs=self.outputs,
            method_name=signature_constants.PREDICT_METHOD_NAME)

        signature_map = {signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: 
                         signature}

        model_builder = saved_model_builder.SavedModelBuilder(model_dir)
        model_builder.add_meta_graph_and_variables(session,
            tags=[tag_constants.SERVING],
            signature_def_map=signature_map,
            clear_devices=True)
        model_builder.save() 

def predict_signature_def(inputs, outputs):
  """Creates prediction signature from given inputs and outputs.

  Args:
    inputs: dict of string to `Tensor`.
    outputs: dict of string to `Tensor`.

  Returns:
    A prediction-flavored signature_def.

  Raises:
    ValueError: If inputs or outputs is `None`.
  """
  if inputs is None or not inputs:
    raise ValueError('inputs cannot be None or empty for prediction.')
  if outputs is None:
    raise ValueError('outputs cannot be None or empty for prediction.')

  # If there's only one input or output, we can standardize keys
  if len(inputs) == 1:
    (_, value), = inputs.items()
    inputs = {signature_constants.PREDICT_INPUTS: value}
  if len(outputs) == 1:
    (_, value), = outputs.items()
    outputs = {signature_constants.PREDICT_OUTPUTS: value}

  signature_inputs = {key: utils.build_tensor_info(tensor)
                      for key, tensor in inputs.items()}
  signature_outputs = {key: utils.build_tensor_info(tensor)
                       for key, tensor in outputs.items()}

  signature_def = build_signature_def(
      signature_inputs, signature_outputs,
      signature_constants.PREDICT_METHOD_NAME)

  return signature_def 

def testConvertNamedSignatureToSignatureDef(self):
    signatures_proto = manifest_pb2.Signatures()
    generic_signature = manifest_pb2.GenericSignature()
    generic_signature.map["input_key"].CopyFrom(
        manifest_pb2.TensorBinding(tensor_name="input"))
    signatures_proto.named_signatures[
        signature_constants.PREDICT_INPUTS].generic_signature.CopyFrom(
            generic_signature)

    generic_signature = manifest_pb2.GenericSignature()
    generic_signature.map["output_key"].CopyFrom(
        manifest_pb2.TensorBinding(tensor_name="output"))
    signatures_proto.named_signatures[
        signature_constants.PREDICT_OUTPUTS].generic_signature.CopyFrom(
            generic_signature)
    signature_def = bundle_shim._convert_named_signatures_to_signature_def(
        signatures_proto)
    self.assertEqual(signature_def.method_name,
                     signature_constants.PREDICT_METHOD_NAME)
    self.assertEqual(len(signature_def.inputs), 1)
    self.assertEqual(len(signature_def.outputs), 1)
    self.assertProtoEquals(
        signature_def.inputs["input_key"],
        meta_graph_pb2.TensorInfo(name="input"))
    self.assertProtoEquals(
        signature_def.outputs["output_key"],
        meta_graph_pb2.TensorInfo(name="output")) 

def predict_signature_def(inputs, outputs):
  """Creates prediction signature from given inputs and outputs.

  Args:
    inputs: dict of string to `Tensor`.
    outputs: dict of string to `Tensor`.

  Returns:
    A prediction-flavored signature_def.

  Raises:
    ValueError: If inputs or outputs is `None`.
  """
  if inputs is None or not inputs:
    raise ValueError('inputs cannot be None or empty for prediction.')
  if outputs is None:
    raise ValueError('outputs cannot be None or empty for prediction.')

  # If there's only one input or output, we can standardize keys
  if len(inputs) == 1:
    (_, value), = inputs.items()
    inputs = {signature_constants.PREDICT_INPUTS: value}
  if len(outputs) == 1:
    (_, value), = outputs.items()
    outputs = {signature_constants.PREDICT_OUTPUTS: value}

  signature_inputs = {key: utils.build_tensor_info(tensor)
                      for key, tensor in inputs.items()}
  signature_outputs = {key: utils.build_tensor_info(tensor)
                       for key, tensor in outputs.items()}

  signature_def = build_signature_def(
      signature_inputs, signature_outputs,
      signature_constants.PREDICT_METHOD_NAME)

  return signature_def 

def testConvertNamedSignatureToSignatureDef(self):
    signatures_proto = manifest_pb2.Signatures()
    generic_signature = manifest_pb2.GenericSignature()
    generic_signature.map["input_key"].CopyFrom(
        manifest_pb2.TensorBinding(tensor_name="input"))
    signatures_proto.named_signatures[
        signature_constants.PREDICT_INPUTS].generic_signature.CopyFrom(
            generic_signature)

    generic_signature = manifest_pb2.GenericSignature()
    generic_signature.map["output_key"].CopyFrom(
        manifest_pb2.TensorBinding(tensor_name="output"))
    signatures_proto.named_signatures[
        signature_constants.PREDICT_OUTPUTS].generic_signature.CopyFrom(
            generic_signature)
    signature_def = bundle_shim._convert_named_signatures_to_signature_def(
        signatures_proto)
    self.assertEqual(signature_def.method_name,
                     signature_constants.PREDICT_METHOD_NAME)
    self.assertEqual(len(signature_def.inputs), 1)
    self.assertEqual(len(signature_def.outputs), 1)
    self.assertProtoEquals(
        signature_def.inputs["input_key"],
        meta_graph_pb2.TensorInfo(name="input"))
    self.assertProtoEquals(
        signature_def.outputs["output_key"],
        meta_graph_pb2.TensorInfo(name="output")) 

def testPredictionSignatureDef(self):
    input1 = constant_op.constant("a", name="input-1")
    input2 = constant_op.constant("b", name="input-2")
    output1 = constant_op.constant("c", name="output-1")
    output2 = constant_op.constant("d", name="output-2")
    signature_def = signature_def_utils.predict_signature_def({
        "input-1": input1,
        "input-2": input2
    }, {"output-1": output1,
        "output-2": output2})

    self.assertEqual(signature_constants.PREDICT_METHOD_NAME,
                     signature_def.method_name)

    # Check inputs in signature def.
    self.assertEqual(2, len(signature_def.inputs))
    input1_tensor_info_actual = (signature_def.inputs["input-1"])
    self.assertEqual("input-1:0", input1_tensor_info_actual.name)
    self.assertEqual(types_pb2.DT_STRING, input1_tensor_info_actual.dtype)
    self.assertEqual(0, len(input1_tensor_info_actual.tensor_shape.dim))
    input2_tensor_info_actual = (signature_def.inputs["input-2"])
    self.assertEqual("input-2:0", input2_tensor_info_actual.name)
    self.assertEqual(types_pb2.DT_STRING, input2_tensor_info_actual.dtype)
    self.assertEqual(0, len(input2_tensor_info_actual.tensor_shape.dim))

    # Check outputs in signature def.
    self.assertEqual(2, len(signature_def.outputs))
    output1_tensor_info_actual = (signature_def.outputs["output-1"])
    self.assertEqual("output-1:0", output1_tensor_info_actual.name)
    self.assertEqual(types_pb2.DT_STRING, output1_tensor_info_actual.dtype)
    self.assertEqual(0, len(output1_tensor_info_actual.tensor_shape.dim))
    output2_tensor_info_actual = (signature_def.outputs["output-2"])
    self.assertEqual("output-2:0", output2_tensor_info_actual.name)
    self.assertEqual(types_pb2.DT_STRING, output2_tensor_info_actual.dtype)
    self.assertEqual(0, len(output2_tensor_info_actual.tensor_shape.dim)) 

def _convert_named_signatures_to_signature_def(signatures):
  """Convert named signatures to object of type SignatureDef.

  Args:
    signatures: object of type manifest_pb2.Signatures()

  Returns:
    object of type SignatureDef which contains a converted version of named
    signatures from input signatures object

  Raises:
    RuntimeError: if input and output named signatures are not of type
    GenericSignature
  """
  signature_def = meta_graph_pb2.SignatureDef()
  input_signature = signatures.named_signatures[
      signature_constants.PREDICT_INPUTS]
  output_signature = signatures.named_signatures[
      signature_constants.PREDICT_OUTPUTS]
  # TODO(pdudnik): what if there are other signatures? Mimic cr/140900781 once
  # it is submitted.
  if (input_signature.WhichOneof("type") != "generic_signature" or
      output_signature.WhichOneof("type") != "generic_signature"):
    raise RuntimeError("Named input and output signatures can only be "
                       "up-converted if they are generic signature. "
                       "Input signature type is %s, output signature type is "
                       "%s" % (input_signature.WhichOneof("type"),
                               output_signature.WhichOneof("type")))

  signature_def.method_name = signature_constants.PREDICT_METHOD_NAME
  for key, val in input_signature.generic_signature.map.items():
    _add_input_to_signature_def(val.tensor_name, key, signature_def)
  for key, val in output_signature.generic_signature.map.items():
    _add_output_to_signature_def(val.tensor_name, key, signature_def)
  return signature_def 

def predict_signature_def(inputs, outputs):
  """Creates prediction signature from given inputs and outputs.

  Args:
    inputs: dict of string to `Tensor`.
    outputs: dict of string to `Tensor`.

  Returns:
    A prediction-flavored signature_def.

  Raises:
    ValueError: If inputs or outputs is `None`.
  """
  if inputs is None or not inputs:
    raise ValueError('inputs cannot be None or empty for prediction.')
  if outputs is None:
    raise ValueError('outputs cannot be None or empty for prediction.')

  signature_inputs = {key: utils.build_tensor_info(tensor)
                      for key, tensor in inputs.items()}
  signature_outputs = {key: utils.build_tensor_info(tensor)
                       for key, tensor in outputs.items()}

  signature_def = build_signature_def(
      signature_inputs, signature_outputs,
      signature_constants.PREDICT_METHOD_NAME)

  return signature_def 

def _write_saved_model(saved_model_path,
                       frozen_graph_def,
                       inputs,
                       outputs):
  """Writes SavedModel to disk.

  If checkpoint_path is not None bakes the weights into the graph thereby
  eliminating the need of checkpoint files during inference. If the model
  was trained with moving averages, setting use_moving_averages to true
  restores the moving averages, otherwise the original set of variables
  is restored.

  Args:
    saved_model_path: Path to write SavedModel.
    frozen_graph_def: tf.GraphDef holding frozen graph.
    inputs: The input image tensor to use for detection.
    outputs: A tensor dictionary containing the outputs of a DetectionModel.
  """
  with tf.Graph().as_default():
    with session.Session() as sess:

      tf.import_graph_def(frozen_graph_def, name='')

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

      tensor_info_inputs = {
          'inputs': tf.saved_model.utils.build_tensor_info(inputs)}
      tensor_info_outputs = {}
      for k, v in outputs.items():
        tensor_info_outputs[k] = tf.saved_model.utils.build_tensor_info(v)

      detection_signature = (
          tf.saved_model.signature_def_utils.build_signature_def(
              inputs=tensor_info_inputs,
              outputs=tensor_info_outputs,
              method_name=signature_constants.PREDICT_METHOD_NAME))

      builder.add_meta_graph_and_variables(
          sess, [tf.saved_model.tag_constants.SERVING],
          signature_def_map={
              signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
                  detection_signature,
          },
      )
      builder.save() 

def _write_saved_model(saved_model_path,
                       frozen_graph_def,
                       inputs,
                       outputs):
  """Writes SavedModel to disk.

  If checkpoint_path is not None bakes the weights into the graph thereby
  eliminating the need of checkpoint files during inference. If the model
  was trained with moving averages, setting use_moving_averages to true
  restores the moving averages, otherwise the original set of variables
  is restored.

  Args:
    saved_model_path: Path to write SavedModel.
    frozen_graph_def: tf.GraphDef holding frozen graph.
    inputs: The input image tensor to use for detection.
    outputs: A tensor dictionary containing the outputs of a DetectionModel.
  """
  with tf.Graph().as_default():
    with session.Session() as sess:

      tf.import_graph_def(frozen_graph_def, name='')

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

      tensor_info_inputs = {
          'inputs': tf.saved_model.utils.build_tensor_info(inputs)}
      tensor_info_outputs = {}
      for k, v in outputs.items():
        tensor_info_outputs[k] = tf.saved_model.utils.build_tensor_info(v)

      detection_signature = (
          tf.saved_model.signature_def_utils.build_signature_def(
              inputs=tensor_info_inputs,
              outputs=tensor_info_outputs,
              method_name=signature_constants.PREDICT_METHOD_NAME))

      builder.add_meta_graph_and_variables(
          sess, [tf.saved_model.tag_constants.SERVING],
          signature_def_map={
              signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
                  detection_signature,
          },
      )
      builder.save() 

def write_saved_model(saved_model_path,
                      frozen_graph_def,
                      inputs,
                      outputs):
  """Writes SavedModel to disk.

  If checkpoint_path is not None bakes the weights into the graph thereby
  eliminating the need of checkpoint files during inference. If the model
  was trained with moving averages, setting use_moving_averages to true
  restores the moving averages, otherwise the original set of variables
  is restored.

  Args:
    saved_model_path: Path to write SavedModel.
    frozen_graph_def: tf.GraphDef holding frozen graph.
    inputs: The input image tensor to use for detection.
    outputs: A tensor dictionary containing the outputs of a DetectionModel.
  """
  with tf.Graph().as_default():
    with session.Session() as sess:

      tf.import_graph_def(frozen_graph_def, name='')

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

      tensor_info_inputs = {
          'inputs': tf.saved_model.utils.build_tensor_info(inputs)}
      tensor_info_outputs = {}
      for k, v in outputs.items():
        tensor_info_outputs[k] = tf.saved_model.utils.build_tensor_info(v)

      detection_signature = (
          tf.saved_model.signature_def_utils.build_signature_def(
              inputs=tensor_info_inputs,
              outputs=tensor_info_outputs,
              method_name=signature_constants.PREDICT_METHOD_NAME))

      builder.add_meta_graph_and_variables(
          sess, [tf.saved_model.tag_constants.SERVING],
          signature_def_map={
              signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
                  detection_signature,
          },
      )
      builder.save() 

def _write_saved_model(saved_model_path,
                       frozen_graph_def,
                       inputs,
                       outputs):
    """Writes SavedModel to disk.

    If checkpoint_path is not None bakes the weights into the graph thereby
    eliminating the need of checkpoint files during inference. If the model
    was trained with moving averages, setting use_moving_averages to true
    restores the moving averages, otherwise the original set of variables
    is restored.

    Args:
      saved_model_path: Path to write SavedModel.
      frozen_graph_def: tf.GraphDef holding frozen graph.
      inputs: The input image tensor to use for detection.
      outputs: A tensor dictionary containing the outputs of a DetectionModel.
    """
    with tf.Graph().as_default():
        with session.Session() as sess:

            tf.import_graph_def(frozen_graph_def, name='')

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

            tensor_info_inputs = {
                'inputs': tf.saved_model.utils.build_tensor_info(inputs)}
            tensor_info_outputs = {}
            for k, v in outputs.items():
                tensor_info_outputs[k] = tf.saved_model.utils.build_tensor_info(
                    v)

            detection_signature = (
                tf.saved_model.signature_def_utils.build_signature_def(
                    inputs=tensor_info_inputs,
                    outputs=tensor_info_outputs,
                    method_name=signature_constants.PREDICT_METHOD_NAME))

            builder.add_meta_graph_and_variables(
                sess, [tf.saved_model.tag_constants.SERVING],
                signature_def_map={
                    signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
                        detection_signature,
                },
            )
            builder.save()