From eacf107c05a5e84b80663d68b5df5769ae12e0c9 Mon Sep 17 00:00:00 2001 From: Stephan Lee Date: Thu, 12 Nov 2020 17:32:42 -0800 Subject: [PATCH 1/3] feat: combine graph by prefixing with unique name Previously, graph plugin combined multiple graphs traced by creating one monolith of a GraphDef. In doing so, we checked whether, for example, node names are unique to detect a case when our graph vis can result in faulty UI. To alleviate the poor UX, we decided, instead, to duplicate all nodes in one gaint GraphDef container prefixing all names. While this creates some bloat, (1) users should see the confusing error less and (2) combined graphs make it very clear that we have traced multiple graphs. --- tensorboard/plugins/graph/BUILD | 4 +- tensorboard/plugins/graph/graph_util.py | 212 +++---- tensorboard/plugins/graph/graph_util_test.py | 604 +++++-------------- tensorboard/plugins/graph/graphs_plugin.py | 12 +- 4 files changed, 245 insertions(+), 587 deletions(-) diff --git a/tensorboard/plugins/graph/BUILD b/tensorboard/plugins/graph/BUILD index 8f2e5a93feb..79aa4a6ffbe 100644 --- a/tensorboard/plugins/graph/BUILD +++ b/tensorboard/plugins/graph/BUILD @@ -123,6 +123,9 @@ py_library( srcs = ["graph_util.py"], srcs_version = "PY2AND3", visibility = ["//visibility:private"], + deps = [ + "//tensorboard/compat/proto:protos_all_py_pb2", + ], ) py_test( @@ -136,7 +139,6 @@ py_test( "//tensorboard:expect_tensorflow_installed", "//tensorboard/compat/proto:protos_all_py_pb2", "@com_google_protobuf//:protobuf_python", - "@org_pythonhosted_six", ], ) diff --git a/tensorboard/plugins/graph/graph_util.py b/tensorboard/plugins/graph/graph_util.py index 0a3d20ef458..b756d359023 100644 --- a/tensorboard/plugins/graph/graph_util.py +++ b/tensorboard/plugins/graph/graph_util.py @@ -14,152 +14,110 @@ # ============================================================================== """Utilities for graph plugin.""" +from tensorboard.compat.proto import graph_pb2 -class _ProtoListDuplicateKeyError(Exception): - pass +def _prefixed_op_name(prefix, op_name): + return "%s/%s" % (prefix, op_name) -class _SameKeyDiffContentError(Exception): - pass +def _prefixed_func_name(prefix, func_name): + # TODO(stephanwlee): add business logic to strip "__inference_". + return "%s_%s" % (prefix, func_name) -def _safe_copy_proto_list_values(dst_proto_list, src_proto_list, get_key): - """Safely merge values from `src_proto_list` into `dst_proto_list`. - Each element in `dst_proto_list` must be mapped by `get_key` to a key - value that is unique within that list; likewise for `src_proto_list`. - If an element of `src_proto_list` has the same key as an existing - element in `dst_proto_list`, then the elements must also be equal. +def _prepend_names(prefix, orig_graph_def): + mut_graph_def = graph_pb2.GraphDef() + for node in orig_graph_def.node: + new_node = mut_graph_def.node.add() + new_node.CopyFrom(node) + new_node.name = _prefixed_op_name(prefix, node.name) + new_node.input[:] = [ + _prefixed_op_name(prefix, input_name) for input_name in node.input + ] - Args: - dst_proto_list: A `RepeatedCompositeContainer` or - `RepeatedScalarContainer` into which values should be copied. - src_proto_list: A container holding the same kind of values as in - `dst_proto_list` from which values should be copied. - get_key: A function that takes an element of `dst_proto_list` or - `src_proto_list` and returns a key, such that if two elements have - the same key then it is required that they be deep-equal. For - instance, if `dst_proto_list` is a list of nodes, then `get_key` - might be `lambda node: node.name` to indicate that if two nodes - have the same name then they must be the same node. All keys must - be hashable. + # Remap tf.function method name in the PartitionedCall. 'f' is short for + # function. + if new_node.op == "PartitionedCall" and new_node.attr["f"]: + + new_node.attr["f"].func.name = _prefixed_func_name( + prefix, new_node.attr["f"].func.name, + ) + + for func in orig_graph_def.library.function: + new_func = mut_graph_def.library.function.add() + new_func.CopyFrom(func) + # Not creating a structure out of factored out function. They already + # create an awkward hierarchy and one for each graph. + new_func.signature.name = _prefixed_func_name( + prefix, new_func.signature.name + ) + + for gradient in orig_graph_def.library.gradient: + new_gradient = mut_graph_def.library.gradient.add() + new_gradient.CopyFrom(gradient) + new_gradient.function_name = _prefixed_func_name( + prefix, new_gradient.function_name, + ) + new_gradient.gradient_func = _prefixed_func_name( + prefix, new_gradient.gradient_func, + ) + + return mut_graph_def - Raises: - _ProtoListDuplicateKeyError: A proto_list contains items with duplicate - keys. - _SameKeyDiffContentError: An item with the same key has different contents. - """ - def _assert_proto_container_unique_keys(proto_list, get_key): - """Asserts proto_list to only contains unique keys. - - Args: - proto_list: A `RepeatedCompositeContainer` or `RepeatedScalarContainer`. - get_key: A function that takes an element of `proto_list` and returns a - hashable key. - - Raises: - _ProtoListDuplicateKeyError: A proto_list contains items with duplicate - keys. - """ - keys = set() - for item in proto_list: - key = get_key(item) - if key in keys: - raise _ProtoListDuplicateKeyError(key) - keys.add(key) - - _assert_proto_container_unique_keys(dst_proto_list, get_key) - _assert_proto_container_unique_keys(src_proto_list, get_key) - - key_to_proto = {} - for proto in dst_proto_list: - key = get_key(proto) - key_to_proto[key] = proto - - for proto in src_proto_list: - key = get_key(proto) - if key in key_to_proto: - if proto != key_to_proto.get(key): - raise _SameKeyDiffContentError(key) - else: - dst_proto_list.add().CopyFrom(proto) - - -def combine_graph_defs(to_proto, from_proto): - """Combines two GraphDefs by adding nodes from from_proto into to_proto. +def merge_graph_defs(graph_defs): + """Merges GraphDefs by adding unique prefix, `graph_{ind}`, to names. All GraphDefs are expected to be of TensorBoard's. - It assumes node names are unique across GraphDefs if contents differ. The - names can be the same if the NodeDef content are exactly the same. + + When collecting graphs using the `tf.summary.trace` API, node names are not + guranteed to be unique. When non-unique names are not considered, it can + lead to graph visualization showing them as one which creates inaccurate + depiction of the flow of the graph (e.g., if there are A -> B -> C and D -> + B -> E, you may see {A, D} -> B -> E). To prevent such graph, we checked + for uniquenss while merging but it resulted in + https://github.com/tensorflow/tensorboard/issues/1929. + + To remedy these issues, we simply "apply name scope" on each graph by + prefixing it with unique name (with a chance of collision) to create + unconnected group of graphs. + + In case there is only one graph def passed, it returns the original + graph_def. In case no graph defs are passed, it returns an empty GraphDef. Args: - to_proto: A destination TensorBoard GraphDef. - from_proto: A TensorBoard GraphDef to copy contents from. + graph_defs: TensorBoard GraphDefs to merge. Returns: - to_proto + TensorBoard GraphDef that merges all graph_defs with unique prefixes. Raises: - ValueError in case any assumption about GraphDef is violated: A - GraphDef should have unique node, function, and gradient function - names. Also, when merging GraphDefs, they should have not have nodes, - functions, or gradient function mappings that share the name but details - do not match. + ValueError in case GraphDef versions mismatch. """ - if from_proto.version != to_proto.version: - raise ValueError("Cannot combine GraphDefs of different versions.") + if len(graph_defs) == 1: + return graph_defs[0] + elif len(graph_defs) == 0: + return graph_pb2.GraphDef() - try: - _safe_copy_proto_list_values( - to_proto.node, from_proto.node, lambda n: n.name - ) - except _ProtoListDuplicateKeyError as exc: - raise ValueError("A GraphDef contains non-unique node names: %s" % exc) - except _SameKeyDiffContentError as exc: - raise ValueError( - ( - "Cannot combine GraphDefs because nodes share a name " - "but contents are different: %s" - ) - % exc - ) - try: - _safe_copy_proto_list_values( - to_proto.library.function, - from_proto.library.function, - lambda n: n.signature.name, - ) - except _ProtoListDuplicateKeyError as exc: - raise ValueError( - "A GraphDef contains non-unique function names: %s" % exc - ) - except _SameKeyDiffContentError as exc: - raise ValueError( - ( - "Cannot combine GraphDefs because functions share a name " - "but are different: %s" - ) - % exc - ) + dst_graph_def = graph_pb2.GraphDef() - try: - _safe_copy_proto_list_values( - to_proto.library.gradient, - from_proto.library.gradient, - lambda g: g.gradient_func, - ) - except _ProtoListDuplicateKeyError as exc: - raise ValueError( - "A GraphDef contains non-unique gradient function names: %s" % exc - ) - except _SameKeyDiffContentError as exc: - raise ValueError( - ( - "Cannot combine GraphDefs because gradients share a gradient_func name " - "but map to different functions: %s" + if graph_defs[0].versions.producer: + dst_graph_def.versions.CopyFrom(graph_defs[0].versions) + + for index, graph_def in enumerate(graph_defs): + if dst_graph_def.versions.producer != graph_def.versions.producer: + raise ValueError("Cannot combine GraphDefs of different versions.") + + mapped_graph_def = _prepend_names("graph_%d" % (index + 1), graph_def) + dst_graph_def.node.extend(mapped_graph_def.node) + if mapped_graph_def.library.function: + dst_graph_def.library.function.extend( + mapped_graph_def.library.function + ) + if mapped_graph_def.library.gradient: + dst_graph_def.library.gradient.extend( + mapped_graph_def.library.gradient ) - % exc - ) - return to_proto + return dst_graph_def diff --git a/tensorboard/plugins/graph/graph_util_test.py b/tensorboard/plugins/graph/graph_util_test.py index 552222a1336..d7073b65780 100644 --- a/tensorboard/plugins/graph/graph_util_test.py +++ b/tensorboard/plugins/graph/graph_util_test.py @@ -12,7 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -import six from google.protobuf import text_format import tensorflow as tf @@ -21,35 +20,43 @@ class GraphUtilTest(tf.test.TestCase): - def test_combine_graph_defs(self): + def test_merge_graph_defs(self): expected_proto = """ node { - name: "X" + name: "graph_1/X" op: "Input" } node { - name: "W" + name: "graph_1/W" op: "Input" } node { - name: "Y" + name: "graph_1/Y" op: "MatMul" - input: "X" - input: "W" + input: "graph_1/X" + input: "graph_1/W" } node { - name: "A" + name: "graph_2/A" op: "Input" } node { - name: "B" + name: "graph_2/B" op: "Input" } node { - name: "C" + name: "graph_2/C" op: "MatMul" - input: "A" - input: "B" + input: "graph_2/A" + input: "graph_2/B" + } + node { + name: "graph_3/A" + op: "Input" + } + node { + name: "graph_3/B" + op: "Input" } versions { producer: 21 @@ -104,31 +111,61 @@ def test_combine_graph_defs(self): graph_def_b, ) + graph_def_c = GraphDef() + text_format.Merge( + """ + node { + name: "A" + op: "Input" + } + node { + name: "B" + op: "Input" + } + versions { + producer: 21 + } + """, + graph_def_c, + ) + self.assertProtoEquals( expected_proto, - graph_util.combine_graph_defs(graph_def_a, graph_def_b), + graph_util.merge_graph_defs( + [graph_def_a, graph_def_b, graph_def_c] + ), ) - def test_combine_graph_defs_name_collided_but_same_content(self): + def test_merge_graph_defs_name_collided_with_same_content(self): expected_proto = """ node { - name: "X" + name: "graph_1/X" op: "Input" } node { - name: "W" + name: "graph_1/W" op: "Input" } node { - name: "Y" + name: "graph_1/Y" op: "MatMul" - input: "X" - input: "W" + input: "graph_1/X" + input: "graph_1/W" + } + node { + name: "graph_2/X" + op: "Input" } node { - name: "A" + name: "graph_2/A" op: "Input" } + node { + name: "graph_2/Y" + op: "MatMul" + input: "graph_2/X" + input: "graph_2/A" + } versions { producer: 21 } @@ -169,107 +206,11 @@ def test_combine_graph_defs_name_collided_but_same_content(self): name: "A" op: "Input" } - versions { - producer: 21 - } - """, - graph_def_b, - ) - - self.assertProtoEquals( - expected_proto, - graph_util.combine_graph_defs(graph_def_a, graph_def_b), - ) - - def test_combine_graph_defs_name_collided_different_content(self): - graph_def_a = GraphDef() - text_format.Merge( - """ - node { - name: "X" - op: "Input" - } - node { - name: "W" - op: "Input" - } node { name: "Y" op: "MatMul" input: "X" - input: "W" - } - versions { - producer: 21 - } - """, - graph_def_a, - ) - - graph_def_b = GraphDef() - text_format.Merge( - """ - node { - name: "X" - op: "Input" - device: "cpu:0" - } - node { - name: "Z" - op: "Input" - } - node { - name: "Q" - op: "MatMul" - input: "X" - input: "Z" - } - versions { - producer: 21 - } - """, - graph_def_b, - ) - - with six.assertRaisesRegex( - self, - ValueError, - ( - "Cannot combine GraphDefs because nodes share a name but " - "contents are different: X" - ), - ): - graph_util.combine_graph_defs(graph_def_a, graph_def_b) - - def test_combine_graph_defs_dst_nodes_duplicate_keys(self): - graph_def_a = GraphDef() - text_format.Merge( - """ - node { - name: "X" - op: "Input" - } - node { - name: "X" - op: "Input" - } - versions { - producer: 21 - } - """, - graph_def_a, - ) - - graph_def_b = GraphDef() - text_format.Merge( - """ - node { - name: "X" - op: "Input" - } - node { - name: "Z" - op: "Input" + input: "A" } versions { producer: 21 @@ -278,60 +219,17 @@ def test_combine_graph_defs_dst_nodes_duplicate_keys(self): graph_def_b, ) - with six.assertRaisesRegex( - self, ValueError, "A GraphDef contains non-unique node names: X" - ): - graph_util.combine_graph_defs(graph_def_a, graph_def_b) - - def test_combine_graph_defs_src_nodes_duplicate_keys(self): - graph_def_a = GraphDef() - text_format.Merge( - """ - node { - name: "X" - op: "Input" - } - node { - name: "Y" - op: "Input" - } - versions { - producer: 21 - } - """, - graph_def_a, - ) - - graph_def_b = GraphDef() - text_format.Merge( - """ - node { - name: "W" - op: "Input" - device: "cpu:0" - } - node { - name: "W" - op: "Input" - } - versions { - producer: 21 - } - """, - graph_def_b, + self.assertProtoEquals( + expected_proto, + graph_util.merge_graph_defs([graph_def_a, graph_def_b]), ) - with six.assertRaisesRegex( - self, ValueError, "A GraphDef contains non-unique node names: W" - ): - graph_util.combine_graph_defs(graph_def_a, graph_def_b) - - def test_combine_graph_defs_function(self): + def test_merge_graph_defs_function(self): expected_proto = """ library { function { signature { - name: "foo" + name: "graph_1_foo" input_arg { name: "x" type: DT_HALF @@ -350,7 +248,26 @@ def test_combine_graph_defs_function(self): } function { signature { - name: "foo_1" + name: "graph_2_foo" + input_arg { + name: "x" + type: DT_INT32 + } + output_arg { + name: "identity" + type: DT_INT32 + } + } + node_def { + name: "add" + op: "Add" + input: "x" + input: "y" + } + } + function { + signature { + name: "graph_2_foo_1" input_arg { name: "x" type: DT_HALF @@ -407,11 +324,11 @@ def test_combine_graph_defs_function(self): name: "foo" input_arg { name: "x" - type: DT_HALF + type: DT_INT32 } output_arg { name: "identity" - type: DT_HALF + type: DT_INT32 } } node_def { @@ -447,64 +364,29 @@ def test_combine_graph_defs_function(self): self.assertProtoEquals( expected_proto, - graph_util.combine_graph_defs(graph_def_a, graph_def_b), + graph_util.merge_graph_defs([graph_def_a, graph_def_b]), ) - def test_combine_graph_defs_function_collison(self): - graph_def_a = GraphDef() + def test_merge_graph_defs_partitioned_call_remap(self): + expected_proto = GraphDef() text_format.Merge( """ - library { - function { - signature { - name: "foo" - input_arg { - name: "x" - type: DT_HALF - } - output_arg { - name: "identity" - type: DT_HALF + node { + name: "graph_1/X" + op: "PartitionedCall" + attr { + key: "f" + value { + func { + name: "graph_1_foo" } } - node_def { - name: "add" - op: "Add" - input: "x" - input: "y" - } } } - """, - graph_def_a, - ) - - graph_def_b = GraphDef() - text_format.Merge( - """ library { function { signature { - name: "foo" - input_arg { - name: "x" - type: DT_HALF - } - output_arg { - name: "identity" - type: DT_HALF - } - } - node_def { - name: "div" - op: "Div" - input: "x" - input: "y" - } - } - function { - signature { - name: "foo_1" + name: "graph_1_foo" input_arg { name: "x" type: DT_HALF @@ -514,109 +396,27 @@ def test_combine_graph_defs_function_collison(self): type: DT_HALF } } - node_def { - name: "add" - op: "Add" - input: "x" - input: "y" - } } } """, - graph_def_b, + expected_proto, ) - with six.assertRaisesRegex( - self, - ValueError, - ( - "Cannot combine GraphDefs because functions share a name but " - "are different: foo" - ), - ): - graph_util.combine_graph_defs(graph_def_a, graph_def_b) - - def test_combine_graph_defs_dst_function_duplicate_keys(self): graph_def_a = GraphDef() text_format.Merge( """ - library { - function { - signature { - name: "foo" - input_arg { - name: "x" - type: DT_HALF - } - output_arg { - name: "identity" - type: DT_HALF - } - } - node_def { - name: "add" - op: "Add" - input: "x" - input: "y" - } - } - function { - signature { - name: "foo" - input_arg { - name: "y" - type: DT_HALF - } - output_arg { - name: "identity" - type: DT_HALF - } - } - } - } - """, - graph_def_a, - ) - - graph_def_b = GraphDef() - text_format.Merge( - """ - library { - function { - signature { - name: "bar" - input_arg { - name: "x" - type: DT_HALF - } - output_arg { - name: "identity" - type: DT_HALF + node { + name: "X" + op: "PartitionedCall" + attr { + key: "f" + value { + func { + name: "foo" } } - node_def { - name: "div" - op: "Div" - input: "x" - input: "y" - } } } - """, - graph_def_b, - ) - - with six.assertRaisesRegex( - self, - ValueError, - ("A GraphDef contains non-unique function names: foo"), - ): - graph_util.combine_graph_defs(graph_def_a, graph_def_b) - - def test_combine_graph_defs_src_function_duplicate_keys(self): - graph_def_a = GraphDef() - text_format.Merge( - """ library { function { signature { @@ -630,70 +430,32 @@ def test_combine_graph_defs_src_function_duplicate_keys(self): type: DT_HALF } } - node_def { - name: "add" - op: "Add" - input: "x" - input: "y" - } } } """, graph_def_a, ) - graph_def_b = GraphDef() - text_format.Merge( - """ - library { - function { - signature { - name: "bar" - input_arg { - name: "x" - type: DT_HALF - } - output_arg { - name: "identity" - type: DT_HALF - } - } - } - function { - signature { - name: "bar" - input_arg { - name: "y" - type: DT_HALF - } - output_arg { - name: "identity" - type: DT_HALF - } - } - } - } - """, - graph_def_b, - ) - with six.assertRaisesRegex( - self, - ValueError, - "A GraphDef contains non-unique function names: bar", - ): - graph_util.combine_graph_defs(graph_def_a, graph_def_b) + self.assertProtoEquals( + expected_proto, + graph_util.merge_graph_defs([graph_def_a, graph_def_b]), + ) - def test_combine_graph_defs_gradient(self): + def test_merge_graph_defs_gradient(self): expected_proto = """ library { gradient { - function_name: "foo" - gradient_func: "foo_grad" + function_name: "graph_1_foo" + gradient_func: "graph_1_foo_grad" + } + gradient { + function_name: "graph_2_foo" + gradient_func: "graph_2_foo_grad" } gradient { - function_name: "bar" - gradient_func: "bar_grad" + function_name: "graph_2_bar" + gradient_func: "graph_2_bar_grad" } } """ @@ -730,126 +492,62 @@ def test_combine_graph_defs_gradient(self): self.assertProtoEquals( expected_proto, - graph_util.combine_graph_defs(graph_def_a, graph_def_b), + graph_util.merge_graph_defs([graph_def_a, graph_def_b]), ) - def test_combine_graph_defs_gradient_collison(self): + def test_merge_graph_defs_mismatch_version(self): graph_def_a = GraphDef() text_format.Merge( """ - library { - gradient { - function_name: "foo" - gradient_func: "foo_grad" - } - } - """, - graph_def_a, - ) - - graph_def_b = GraphDef() - text_format.Merge( - """ - library { - gradient { - function_name: "bar" - gradient_func: "bar_grad" - } - gradient { - function_name: "foo_1" - gradient_func: "foo_grad" - } - } - """, - graph_def_b, - ) - - with six.assertRaisesRegex( - self, - ValueError, - ( - "share a gradient_func name but map to different functions: " - "foo_grad" - ), - ): - graph_util.combine_graph_defs(graph_def_a, graph_def_b) - - def test_combine_graph_defs_dst_gradient_func_non_unique(self): - graph_def_a = GraphDef() - text_format.Merge( - """ - library { - gradient { - function_name: "foo" - gradient_func: "foo_grad" - } - gradient { - function_name: "foo_bar" - gradient_func: "foo_grad" - } - } - """, + node { + name: "A" + op: "Input" + } + versions { + producer: 21 + } + """, graph_def_a, ) graph_def_b = GraphDef() text_format.Merge( """ - library { - gradient { - function_name: "bar" - gradient_func: "bar_grad" - } - } - """, + node { + name: "A" + op: "Input" + } + versions { + producer: 100 + } + """, graph_def_b, ) - with six.assertRaisesRegex( - self, - ValueError, - "A GraphDef contains non-unique gradient function names: foo_grad", + with self.assertRaisesRegex( + ValueError, "Cannot combine GraphDefs of different versions" ): - graph_util.combine_graph_defs(graph_def_a, graph_def_b) + graph_util.merge_graph_defs([graph_def_a, graph_def_b]) - def test_combine_graph_defs_src_gradient_func_non_unique(self): + def test_merge_graph_defs_single_graph_def_no_prefix(self): graph_def_a = GraphDef() text_format.Merge( """ - library { - gradient { - function_name: "foo" - gradient_func: "foo_grad" - } - } - """, + node { + name: "A" + op: "Input" + } + versions { + producer: 21 + } + """, graph_def_a, ) - graph_def_b = GraphDef() - text_format.Merge( - """ - library { - gradient { - function_name: "bar" - gradient_func: "bar_grad" - } - gradient { - function_name: "bar_baz" - gradient_func: "bar_grad" - } - } - """, - graph_def_b, + self.assertProtoEquals( + graph_def_a, graph_util.merge_graph_defs([graph_def_a]), ) - with six.assertRaisesRegex( - self, - ValueError, - "A GraphDef contains non-unique gradient function names: bar_grad", - ): - graph_util.combine_graph_defs(graph_def_a, graph_def_b) - if __name__ == "__main__": tf.test.main() diff --git a/tensorboard/plugins/graph/graphs_plugin.py b/tensorboard/plugins/graph/graphs_plugin.py index f5a1d8bd5be..9c9cb978347 100644 --- a/tensorboard/plugins/graph/graphs_plugin.py +++ b/tensorboard/plugins/graph/graphs_plugin.py @@ -243,12 +243,12 @@ def graph_impl( run_metadata = config_pb2.RunMetadata.FromString( tensor_events[0].tensor_proto.string_val[0] ) - graph = graph_pb2.GraphDef() - - for func_graph in run_metadata.function_graphs: - graph_util.combine_graph_defs( - graph, func_graph.pre_optimization_graph - ) + graph = graph_util.merge_graph_defs( + [ + func_graph.pre_optimization_graph + for func_graph in run_metadata.function_graphs + ] + ) else: graph = self._multiplexer.Graph(run) From 883073ac34bcdf7578b95b19d250fdf28c8e16eb Mon Sep 17 00:00:00 2001 From: Stephan Lee Date: Tue, 17 Nov 2020 16:07:30 -0800 Subject: [PATCH 2/3] Cr address --- tensorboard/plugins/graph/graph_util.py | 49 +++++++++++++------------ 1 file changed, 26 insertions(+), 23 deletions(-) diff --git a/tensorboard/plugins/graph/graph_util.py b/tensorboard/plugins/graph/graph_util.py index b756d359023..2f5f51ff04b 100644 --- a/tensorboard/plugins/graph/graph_util.py +++ b/tensorboard/plugins/graph/graph_util.py @@ -22,14 +22,26 @@ def _prefixed_op_name(prefix, op_name): def _prefixed_func_name(prefix, func_name): - # TODO(stephanwlee): add business logic to strip "__inference_". + """Returns function name prefixed with `prefix`. + + For function libraries, which are often created out of autographed Python + function, are factored out in the graph vis. They are grouped under a + function name which often has a shape of + `__inference_[py_func_name]_[numeric_suffix]`. + + While it does not have some unique information about which graph it is from, + creating another wrapping structure with graph prefix and "/" is less than + ideal so we join the prefix and func_name using underscore. + + TODO(stephanwlee): add business logic to strip "__inference_" for more user + friendlier name + """ return "%s_%s" % (prefix, func_name) -def _prepend_names(prefix, orig_graph_def): - mut_graph_def = graph_pb2.GraphDef() - for node in orig_graph_def.node: - new_node = mut_graph_def.node.add() +def _add_with_prepended_names(prefix, graph_to_add, destination_graph): + for node in graph_to_add.node: + new_node = destination_graph.node.add() new_node.CopyFrom(node) new_node.name = _prefixed_op_name(prefix, node.name) new_node.input[:] = [ @@ -44,17 +56,15 @@ def _prepend_names(prefix, orig_graph_def): prefix, new_node.attr["f"].func.name, ) - for func in orig_graph_def.library.function: - new_func = mut_graph_def.library.function.add() + for func in graph_to_add.library.function: + new_func = destination_graph.library.function.add() new_func.CopyFrom(func) - # Not creating a structure out of factored out function. They already - # create an awkward hierarchy and one for each graph. new_func.signature.name = _prefixed_func_name( prefix, new_func.signature.name ) - for gradient in orig_graph_def.library.gradient: - new_gradient = mut_graph_def.library.gradient.add() + for gradient in graph_to_add.library.gradient: + new_gradient = destination_graph.library.gradient.add() new_gradient.CopyFrom(gradient) new_gradient.function_name = _prefixed_func_name( prefix, new_gradient.function_name, @@ -63,8 +73,6 @@ def _prepend_names(prefix, orig_graph_def): prefix, new_gradient.gradient_func, ) - return mut_graph_def - def merge_graph_defs(graph_defs): """Merges GraphDefs by adding unique prefix, `graph_{ind}`, to names. @@ -109,15 +117,10 @@ def merge_graph_defs(graph_defs): if dst_graph_def.versions.producer != graph_def.versions.producer: raise ValueError("Cannot combine GraphDefs of different versions.") - mapped_graph_def = _prepend_names("graph_%d" % (index + 1), graph_def) - dst_graph_def.node.extend(mapped_graph_def.node) - if mapped_graph_def.library.function: - dst_graph_def.library.function.extend( - mapped_graph_def.library.function - ) - if mapped_graph_def.library.gradient: - dst_graph_def.library.gradient.extend( - mapped_graph_def.library.gradient - ) + _add_with_prepended_names( + "graph_%d" % (index + 1), + graph_def, + dst_graph_def, + ) return dst_graph_def From d6df7b607c32e6ad0917c8b0b53756f9ac483012 Mon Sep 17 00:00:00 2001 From: Stephan Lee Date: Tue, 17 Nov 2020 16:12:55 -0800 Subject: [PATCH 3/3] whoops --- tensorboard/plugins/graph/graph_util.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tensorboard/plugins/graph/graph_util.py b/tensorboard/plugins/graph/graph_util.py index 2f5f51ff04b..879554b4d93 100644 --- a/tensorboard/plugins/graph/graph_util.py +++ b/tensorboard/plugins/graph/graph_util.py @@ -118,9 +118,7 @@ def merge_graph_defs(graph_defs): raise ValueError("Cannot combine GraphDefs of different versions.") _add_with_prepended_names( - "graph_%d" % (index + 1), - graph_def, - dst_graph_def, + "graph_%d" % (index + 1), graph_def, dst_graph_def, ) return dst_graph_def