Python saved_model_pb2.SavedModel方法代码示例

# 需要导入模块: from tensorflow.core.protobuf import saved_model_pb2 [as 别名]
# 或者: from tensorflow.core.protobuf.saved_model_pb2 import SavedModel [as 别名]
def __init__(self, export_dir):
    self._saved_model = saved_model_pb2.SavedModel()
    self._saved_model.saved_model_schema_version = (
        constants.SAVED_MODEL_SCHEMA_VERSION)

    self._export_dir = export_dir
    if file_io.file_exists(export_dir):
      raise AssertionError(
          "Export directory already exists. Please specify a different export "
          "directory: %s" % export_dir)

    file_io.recursive_create_dir(self._export_dir)

    # Boolean to track whether variables and assets corresponding to the
    # SavedModel have been saved. Specifically, the first meta graph to be added
    # MUST use the add_meta_graph_and_variables() API. Subsequent add operations
    # on the SavedModel MUST use the add_meta_graph() API which does not save
    # weights.
    self._has_saved_variables = False 

# 需要导入模块: from tensorflow.core.protobuf import saved_model_pb2 [as 别名]
# 或者: from tensorflow.core.protobuf.saved_model_pb2 import SavedModel [as 别名]
def _tag_and_add_meta_graph(self, meta_graph_def, tags, signature_def_map):
    """Tags the meta graph def and adds it to the SavedModel.

    Tags the meta graph def with the supplied tags, adds signature defs to it if
    provided and appends the meta graph def to the SavedModel proto.

    Args:
      meta_graph_def: The meta graph def to add to the SavedModel.
      tags: The set of tags to annotate the meta graph def with.
      signature_def_map: The map of signature defs to be added to the meta graph
          def.
    """
    for tag in tags:
      meta_graph_def.meta_info_def.tags.append(tag)

    if signature_def_map is not None:
      for key in signature_def_map:
        meta_graph_def.signature_def[key].CopyFrom(signature_def_map[key])

    proto_meta_graph_def = self._saved_model.meta_graphs.add()
    proto_meta_graph_def.CopyFrom(meta_graph_def) 

# 需要导入模块: from tensorflow.core.protobuf import saved_model_pb2 [as 别名]
# 或者: from tensorflow.core.protobuf.saved_model_pb2 import SavedModel [as 别名]
def _get_main_op_tensor(meta_graph_def_to_load):
  """Gets the main op tensor, if one exists.

  Args:
    meta_graph_def_to_load: The meta graph def from the SavedModel to be loaded.

  Returns:
    The main op tensor, if it exists and `None` otherwise.

  Raises:
    RuntimeError: If the collection def corresponding to the main op key has
        other than exactly one tensor.
  """
  collection_def = meta_graph_def_to_load.collection_def
  main_op_tensor = None
  if constants.MAIN_OP_KEY in collection_def:
    main_ops = collection_def[constants.MAIN_OP_KEY].node_list.value
    if len(main_ops) != 1:
      raise RuntimeError("Expected exactly one SavedModel main op.")
    main_op_tensor = ops.get_collection(constants.MAIN_OP_KEY)[0]
  return main_op_tensor 

# 需要导入模块: from tensorflow.core.protobuf import saved_model_pb2 [as 别名]
# 或者: from tensorflow.core.protobuf.saved_model_pb2 import SavedModel [as 别名]
def _get_legacy_init_op_tensor(meta_graph_def_to_load):
  """Gets the legacy init op tensor, if one exists.

  Args:
    meta_graph_def_to_load: The meta graph def from the SavedModel to be loaded.

  Returns:
    The legacy init op tensor, if it exists and `None` otherwise.

  Raises:
    RuntimeError: If the collection def corresponding to the legacy init op key
        has other than exactly one tensor.
  """
  collection_def = meta_graph_def_to_load.collection_def
  legacy_init_op_tensor = None
  if constants.LEGACY_INIT_OP_KEY in collection_def:
    legacy_init_ops = collection_def[
        constants.LEGACY_INIT_OP_KEY].node_list.value
    if len(legacy_init_ops) != 1:
      raise RuntimeError("Expected exactly one legacy serving init op.")
    legacy_init_op_tensor = ops.get_collection(constants.LEGACY_INIT_OP_KEY)[0]
  return legacy_init_op_tensor 

# 需要导入模块: from tensorflow.core.protobuf import saved_model_pb2 [as 别名]
# 或者: from tensorflow.core.protobuf.saved_model_pb2 import SavedModel [as 别名]
def maybe_saved_model_directory(export_dir):
  """Checks whether the provided export directory could contain a SavedModel.

  Note that the method does not load any data by itself. If the method returns
  `false`, the export directory definitely does not contain a SavedModel. If the
  method returns `true`, the export directory may contain a SavedModel but
  provides no guarantee that it can be loaded.

  Args:
    export_dir: Absolute string path to possible export location. For example,
                '/my/foo/model'.

  Returns:
    True if the export directory contains SavedModel files, False otherwise.
  """
  txt_path = os.path.join(export_dir, constants.SAVED_MODEL_FILENAME_PBTXT)
  pb_path = os.path.join(export_dir, constants.SAVED_MODEL_FILENAME_PB)
  return file_io.file_exists(txt_path) or file_io.file_exists(pb_path) 

# 需要导入模块: from tensorflow.core.protobuf import saved_model_pb2 [as 别名]
# 或者: from tensorflow.core.protobuf.saved_model_pb2 import SavedModel [as 别名]
def __init__(self, export_dir):
    self._saved_model = saved_model_pb2.SavedModel()
    self._saved_model.saved_model_schema_version = (
        constants.SAVED_MODEL_SCHEMA_VERSION)

    self._export_dir = export_dir
    if file_io.file_exists(export_dir):
      raise AssertionError(
          "Export directory already exists. Please specify a different export "
          "directory.")

    file_io.recursive_create_dir(self._export_dir)

    # Boolean to track whether variables and assets corresponding to the
    # SavedModel have been saved. Specifically, the first meta graph to be added
    # MUST use the add_meta_graph_and_variables() API. Subsequent add operations
    # on the SavedModel MUST use the add_meta_graph() API which does not save
    # weights.
    self._has_saved_variables = False 

# 需要导入模块: from tensorflow.core.protobuf import saved_model_pb2 [as 别名]
# 或者: from tensorflow.core.protobuf.saved_model_pb2 import SavedModel [as 别名]
def _get_legacy_init_op_tensor(meta_graph_def_to_load):
  """Gets the legacy init op tensor, if one exists.

  Args:
    meta_graph_def_to_load: The meta graph def from the SavedModel to be loaded.

  Returns:
    The legacy init op tensor, if it exists and `None` otherwise.

  Raises:
    RuntimeError: If the collection def corresponding to the legacy init op key
        has other than exactly one tensor.
  """
  collection_def = meta_graph_def_to_load.collection_def
  legacy_init_op_tensor = None
  if constants.LEGACY_INIT_OP_KEY in collection_def:
    legacy_init_ops = collection_def[
        constants.LEGACY_INIT_OP_KEY].node_list.value
    if len(legacy_init_ops) != 1:
      raise RuntimeError("Expected exactly one legacy serving init op.")
    legacy_init_op_tensor = tf.get_collection(constants.LEGACY_INIT_OP_KEY)[0]
  return legacy_init_op_tensor 

# 需要导入模块: from tensorflow.core.protobuf import saved_model_pb2 [as 别名]
# 或者: from tensorflow.core.protobuf.saved_model_pb2 import SavedModel [as 别名]
def import_to_tensorboard(saved_model, output_dir):
  """View an imported saved_model.pb as a graph in Tensorboard.

  Args:
    saved_model: The location of the saved_model.pb to visualize.
    output_dir: The location for the Tensorboard log to begin visualization from.

  Usage:
    Call this function with your model location and desired log directory.
    Launch Tensorboard by pointing it to the log directory.
    View your imported `.pb` model as a graph.
  """
  with open(saved_model, "rb") as f:
    sm = saved_model_pb2.SavedModel()
    sm.ParseFromString(f.read())
    if 1 != len(sm.meta_graphs):
      print('More than one graph found. Not sure which to write')
      sys.exit(1)
    graph_def = sm.meta_graphs[0].graph_def

    pb_visual_writer = summary.FileWriter(output_dir)
    pb_visual_writer.add_graph(None, graph_def=graph_def)
    print("Model Imported. Visualize by running: "
          "tensorboard --logdir={}".format(output_dir)) 

# 需要导入模块: from tensorflow.core.protobuf import saved_model_pb2 [as 别名]
# 或者: from tensorflow.core.protobuf.saved_model_pb2 import SavedModel [as 别名]
def export(self, path, variables_saver=None):
    """Exports to SavedModel directory.

    Args:
      path: path where to export the SavedModel to.
      variables_saver: lambda that receives a directory path where to
        export checkpoints of variables.
    """
    # Operate on a copy of self._proto since it needs to be modified.
    proto = saved_model_pb2.SavedModel()
    proto.CopyFrom(self._proto)
    assets_map = _make_assets_key_collection(proto, path)

    self._save_all_assets(path, assets_map)
    self._save_variables(path, variables_saver)
    self._save_proto(path, proto) 

# 需要导入模块: from tensorflow.core.protobuf import saved_model_pb2 [as 别名]
# 或者: from tensorflow.core.protobuf.saved_model_pb2 import SavedModel [as 别名]
def _maybe_add_legacy_init_op(self, legacy_init_op=None):
    """Add legacy init op to the SavedModel.

    Args:
      legacy_init_op: Optional legacy init op to support backward compatibility.

    Raises:
      TypeError if legacy init op is not of type `Operation`.
      AssertionError if the graph already contains one or more legacy init ops.
    """
    if legacy_init_op is not None:
      if not isinstance(legacy_init_op, ops.Operation):
        raise TypeError("legacy_init_op needs to be an Operation: %r" %
                        legacy_init_op)
      if ops.get_collection(constants.LEGACY_INIT_OP_KEY):
        raise AssertionError(
            "graph already contains one or more legacy init ops under the "
            "collection {}.".format(constants.LEGACY_INIT_OP_KEY))
      ops.add_to_collection(constants.LEGACY_INIT_OP_KEY, legacy_init_op) 

# 需要导入模块: from tensorflow.core.protobuf import saved_model_pb2 [as 别名]
# 或者: from tensorflow.core.protobuf.saved_model_pb2 import SavedModel [as 别名]
def _maybe_add_legacy_init_op(self, legacy_init_op=None):
    """Add legacy init op to the SavedModel.

    Args:
      legacy_init_op: Optional legacy init op to support backward compatibility.

    Raises:
      TypeError if legacy init op is not of type `Operation`.
    """
    if legacy_init_op is not None:
      if not isinstance(legacy_init_op, ops.Operation):
        raise TypeError("legacy_init_op needs to be an Operation: %r" %
                        legacy_init_op)
      ops.add_to_collection(constants.LEGACY_INIT_OP_KEY, legacy_init_op) 

# 需要导入模块: from tensorflow.core.protobuf import saved_model_pb2 [as 别名]
# 或者: from tensorflow.core.protobuf.saved_model_pb2 import SavedModel [as 别名]
def _add_main_op(self, main_op):
    """Add main op to the SavedModel.

    Args:
      main_op: Main op to run as part of graph initialization.

    Raises:
      TypeError if main op is not of type `Operation`.
    """
    if main_op is not None:
      if not isinstance(main_op, ops.Operation):
        raise TypeError("main_op needs to be an Operation: %r" % main_op)
      ops.add_to_collection(constants.MAIN_OP_KEY, main_op) 

# 需要导入模块: from tensorflow.core.protobuf import saved_model_pb2 [as 别名]
# 或者: from tensorflow.core.protobuf.saved_model_pb2 import SavedModel [as 别名]
def _get_asset_tensors(export_dir, meta_graph_def_to_load):
  """Gets the asset tensors, if defined in the meta graph def to load.

  Args:
    export_dir: Directory where the SavedModel is located.
    meta_graph_def_to_load: The meta graph def from the SavedModel to be loaded.

  Returns:
    A dictionary of asset tensors, keyed by the name of the asset tensor. The
    value in the map corresponds to the absolute path of the asset file.
  """
  # Collection-def that may contain the assets key.
  collection_def = meta_graph_def_to_load.collection_def

  asset_tensor_dict = {}
  if constants.ASSETS_KEY in collection_def:
    # Location of the assets for SavedModel.
    assets_directory = os.path.join(
        compat.as_bytes(export_dir),
        compat.as_bytes(constants.ASSETS_DIRECTORY))
    assets_any_proto = collection_def[constants.ASSETS_KEY].any_list.value
    # Process each asset and add it to the asset tensor dictionary.
    for asset_any_proto in assets_any_proto:
      asset_proto = meta_graph_pb2.AssetFileDef()
      asset_any_proto.Unpack(asset_proto)
      asset_tensor_dict[asset_proto.tensor_info.name] = os.path.join(
          compat.as_bytes(assets_directory),
          compat.as_bytes(asset_proto.filename))
  return asset_tensor_dict