refactor

~tensorboardX/graph.py

@@ -0,0 +1,73 @@
+from distutils.version import LooseVersion
+
+from .src.attr_value_pb2 import AttrValue
+from .src.graph_pb2 import GraphDef
+from .src.node_def_pb2 import NodeDef
+from .src.tensor_shape_pb2 import TensorShapeProto
+from .src.versions_pb2 import VersionDef
+
+
+def replace(name, scope):
+    return "/".join([scope[name], name])
+
+
+def parse(graph):
+    import torch
+
+    scope = {}
+    for n in graph.nodes():
+        inputs = [i.uniqueName() for i in n.inputs()]
+        for i in range(1, len(inputs)):
+            scope[inputs[i]] = n.scopeName()
+
+        uname = next(iter(n.outputs())).uniqueName()
+        assert n.scopeName() != "", "{} has empty scope name".format(n)
+        scope[uname] = n.scopeName()
+    if LooseVersion(torch.__version__) >= LooseVersion("0.4"):
+        scope["0"] = "input"
+    else:
+        scope["1"] = "input"
+
+    nodes = []
+    for n in graph.nodes():
+        attrs = {k: n[k] for k in n.attributeNames()}
+        attrs = str(attrs).replace("'", " ")  # singlequote will be escaped by tensorboard
+        if any(i.uniqueName() not in scope.keys() for i in n.inputs()):  # 0.3.1 workaround
+            continue
+        inputs = [replace(i.uniqueName(), scope) for i in n.inputs()]
+        uname = next(iter(n.outputs())).uniqueName()  # FIXME: only first output is considered
+        nodes.append({"name": replace(uname, scope), "op": n.kind(), "inputs": inputs, "attr": attrs})
+
+    for n in graph.inputs():
+        uname = n.uniqueName()
+        if uname not in scope.keys():
+            scope[uname] = "unused"
+        nodes.append({"name": replace(uname, scope), "op": "Parameter", "inputs": [], "attr": str(n.type())})
+
+    return nodes
+
+
+def graph(model, args, verbose=False):
+    import torch
+
+    with torch.onnx.set_training(model, False):
+        trace, _ = torch.jit.trace(model, args)
+    if LooseVersion(torch.__version__) >= LooseVersion("0.4"):
+        torch.onnx._optimize_trace(trace, False)
+    else:
+        torch.onnx._optimize_trace(trace)
+    graph = trace.graph()
+    if verbose:
+        print(graph)
+    list_of_nodes = parse(graph)
+    nodes = []
+    for node in list_of_nodes:
+        nodes.append(
+            NodeDef(
+                name=node["name"],
+                op=node["op"],
+                input=node["inputs"],
+                attr={"lanpa": AttrValue(s=node["attr"].encode(encoding="utf_8"))},
+            )
+        )
+    return GraphDef(node=nodes, versions=VersionDef(producer=22))