def testNoModuleAttachments(self): meta_graph = meta_graph_pb2.MetaGraphDef() with tf.Graph().as_default(): # No calls to attach_bytes. saved_model_lib._export_module_attachments(meta_graph) actual = saved_model_lib.get_attached_bytes_map(meta_graph) self.assertDictEqual({}, actual) # Check there were no unwarranted subscript operations. self.assertFalse(saved_model_lib.ATTACHMENT_COLLECTION_SAVED in meta_graph.collection_def)
def testModuleAttachments(self): meta_graph = meta_graph_pb2.MetaGraphDef() with tf.Graph().as_default(): saved_model_lib.attach_bytes("key1", tf.compat.as_bytes("oops")) saved_model_lib.attach_bytes("key2", tf.compat.as_bytes("value2")) saved_model_lib.attach_bytes("key1", tf.compat.as_bytes("value1")) saved_model_lib._export_module_attachments(meta_graph) actual = saved_model_lib.get_attached_bytes_map(meta_graph) expected = {"key1": tf.compat.as_bytes("value1"), "key2": tf.compat.as_bytes("value2")} self.assertDictEqual(expected, actual)