Python meta_graph_pb2.AssetFileDef方法代码示例

# 需要导入模块: from tensorflow.core.protobuf import meta_graph_pb2 [as 别名]
# 或者: from tensorflow.core.protobuf.meta_graph_pb2 import AssetFileDef [as 别名]
def _validate_asset_collection(self, export_dir, graph_collection_def,
                                 expected_asset_file_name,
                                 expected_asset_file_contents,
                                 expected_asset_tensor_name):
    assets_any = graph_collection_def[constants.ASSETS_KEY].any_list.value
    asset = meta_graph_pb2.AssetFileDef()
    assets_any[0].Unpack(asset)
    assets_path = os.path.join(
        compat.as_bytes(export_dir),
        compat.as_bytes(constants.ASSETS_DIRECTORY),
        compat.as_bytes(expected_asset_file_name))
    actual_asset_contents = file_io.read_file_to_string(assets_path)
    self.assertEqual(expected_asset_file_contents,
                     compat.as_text(actual_asset_contents))
    self.assertEqual(expected_asset_file_name, asset.filename)
    self.assertEqual(expected_asset_tensor_name, asset.tensor_info.name) 

# 需要导入模块: from tensorflow.core.protobuf import meta_graph_pb2 [as 别名]
# 或者: from tensorflow.core.protobuf.meta_graph_pb2 import AssetFileDef [as 别名]
def _add_asset_to_collection(asset_filename, asset_tensor):
  """Builds an asset proto and adds it to the asset collection of the graph.

  Args:
    asset_filename: The filename of the asset to be added.
    asset_tensor: The asset tensor used to populate the tensor info of the
        asset proto.
  """
  asset_proto = meta_graph_pb2.AssetFileDef()
  asset_proto.filename = asset_filename
  asset_proto.tensor_info.name = asset_tensor.name

  asset_any_proto = Any()
  asset_any_proto.Pack(asset_proto)
  ops.add_to_collection(constants.ASSETS_KEY, asset_any_proto) 

# 需要导入模块: from tensorflow.core.protobuf import meta_graph_pb2 [as 别名]
# 或者: from tensorflow.core.protobuf.meta_graph_pb2 import AssetFileDef [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 

# 需要导入模块: from tensorflow.core.protobuf import meta_graph_pb2 [as 别名]
# 或者: from tensorflow.core.protobuf.meta_graph_pb2 import AssetFileDef [as 别名]
def _add_asset_to_collection(self, asset_filename, asset_tensor):
    """Builds an asset proto and adds it to the asset collection of the graph.

    Args:
      asset_filename: The filename of the asset to be added.
      asset_tensor: The asset tensor used to populate the tensor info of the
          asset proto.
    """
    asset_proto = meta_graph_pb2.AssetFileDef()
    asset_proto.filename = asset_filename
    asset_proto.tensor_info.name = asset_tensor.name

    asset_any_proto = Any()
    asset_any_proto.Pack(asset_proto)
    ops.add_to_collection(constants.ASSETS_KEY, asset_any_proto) 

# 需要导入模块: from tensorflow.core.protobuf import meta_graph_pb2 [as 别名]
# 或者: from tensorflow.core.protobuf.meta_graph_pb2 import AssetFileDef [as 别名]
def testCreationOfAssetsKeyCollectionIsDeterministic(self):
    tmp_asset_dir = os.path.join(self.get_temp_dir(), "assets")
    tf_v1.gfile.MakeDirs(tmp_asset_dir)
    filenames = [
        os.path.join(tmp_asset_dir, "file%d.txt" % n) for n in range(10)
    ]
    for filename in filenames:
      _write_string_to_file(filename, "I am file %s" % filename)

    with tf.Graph().as_default() as graph:
      assets = [tf.constant(f, name=os.path.basename(f)) for f in filenames]
      for asset in assets:
        graph.add_to_collection(tf_v1.GraphKeys.ASSET_FILEPATHS, asset)
      saved_model_lib.add_signature("default", {}, {"default": assets[0]})

    handler = saved_model_lib.SavedModelHandler()
    handler.add_graph_copy(graph)
    saved_model_proto = copy.deepcopy(handler._proto)
    export_dir = os.path.join(self.get_temp_dir(), "assets_key_test")
    saved_model_lib._make_assets_key_collection(saved_model_proto, export_dir)

    meta_graph = list(saved_model_proto.meta_graphs)[0]
    asset_tensor_names = []
    for asset_any_proto in meta_graph.collection_def[
        tf_v1.saved_model.constants.ASSETS_KEY].any_list.value:
      asset_proto = meta_graph_pb2.AssetFileDef()
      asset_any_proto.Unpack(asset_proto)
      asset_tensor_names.append(asset_proto.tensor_info.name)
    self.assertEqual(asset_tensor_names, sorted(asset_tensor_names)) 

# 需要导入模块: from tensorflow.core.protobuf import meta_graph_pb2 [as 别名]
# 或者: from tensorflow.core.protobuf.meta_graph_pb2 import AssetFileDef [as 别名]
def _merge_assets_key_collection(saved_model_proto, path):
  """Merges the ASSETS_KEY collection into the GraphDefs in saved_model_proto.

  Removes the ASSETS_KEY collection from the GraphDefs in the SavedModel and
  modifies nodes with the assets filenames to point to the assets in `path`.
  After this transformation, the SavedModel GraphDefs can be used without
  feeding asset tensors.

  Args:
    saved_model_proto: SavedModel proto to be modified.
    path: path where the SavedModel is being loaded from.
  """
  for meta_graph in saved_model_proto.meta_graphs:
    node_asset_map = {}
    if tf_v1.saved_model.constants.ASSETS_KEY in meta_graph.collection_def:
      assets_any_proto = meta_graph.collection_def[
          tf_v1.saved_model.constants.ASSETS_KEY].any_list.value
      for asset_any_proto in assets_any_proto:
        asset_proto = meta_graph_pb2.AssetFileDef()
        asset_any_proto.Unpack(asset_proto)
        asset_filename = _get_asset_filename(path, asset_proto.filename)
        node_asset_map[_get_node_name_from_tensor(
            asset_proto.tensor_info.name)] = asset_filename
      del meta_graph.collection_def[tf_v1.saved_model.constants.ASSETS_KEY]

    for node in meta_graph.graph_def.node:
      asset_filepath = node_asset_map.get(node.name)
      if asset_filepath:
        _check_asset_node_def(node)
        node.attr["value"].tensor.string_val[0] = asset_filepath