def get_attached_bytes_map(meta_graph):
"""Returns the dict of ModuleAttachments stored in `meta_graph`.
Args:
meta_graph: A MetaGraphDef, as built by SavedModelHandler.add_graph_copy()
from some graph.
Returns:
A dict, containing the `(key, bytes)` items passed to `attach_bytes()`
when the graph had been built.
Raises:
ValueError: if `meta-graph` is malformed.
"""
result = {}
if ATTACHMENT_COLLECTION_SAVED not in meta_graph.collection_def:
return result
collection_def = meta_graph.collection_def[ATTACHMENT_COLLECTION_SAVED]
if collection_def.WhichOneof("kind") != "bytes_list":
raise ValueError(
"Internal CollectionDef for attached messages has kind %s, "
"expected bytes_list" % collection_def.WhichOneof("kind"))
attachment = module_attachment_pb2.ModuleAttachment()
for value in collection_def.bytes_list.value:
attachment.ParseFromString(value)
result[attachment.key] = attachment.value # Immutable; needs no copy.
return result
def attach_bytes(key, the_bytes):
"""Adds a ModuleAttachment to the current graph.
Args:
key: A string with the unique key of the attachment.
the_bytes: A bytes object with the serialized attachment.
"""
tf_v1.add_to_collection(
_ATTACHMENT_COLLECTION_INTERNAL,
module_attachment_pb2.ModuleAttachment(key=key, value=the_bytes))
