File size: 4,791 Bytes
cf2a15a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
# Copyright 2019 The TensorFlow Authors. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Utilities for graph plugin."""

from tensorboard.compat.proto import graph_pb2


def _prefixed_op_name(prefix, op_name):
    return "%s/%s" % (prefix, op_name)


def _prefixed_func_name(prefix, func_name):
    """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 _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[:] = [
            _prefixed_op_name(prefix, input_name) for input_name in node.input
        ]

        # 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 graph_to_add.library.function:
        new_func = destination_graph.library.function.add()
        new_func.CopyFrom(func)
        new_func.signature.name = _prefixed_func_name(
            prefix, new_func.signature.name
        )

    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,
        )
        new_gradient.gradient_func = _prefixed_func_name(
            prefix,
            new_gradient.gradient_func,
        )


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.

    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:
      graph_defs: TensorBoard GraphDefs to merge.

    Returns:
      TensorBoard GraphDef that merges all graph_defs with unique prefixes.

    Raises:
      ValueError in case GraphDef versions mismatch.
    """
    if len(graph_defs) == 1:
        return graph_defs[0]
    elif len(graph_defs) == 0:
        return graph_pb2.GraphDef()

    dst_graph_def = graph_pb2.GraphDef()

    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.")

        _add_with_prepended_names(
            "graph_%d" % (index + 1),
            graph_def,
            dst_graph_def,
        )

    return dst_graph_def