74 lines
2.4 KiB
Python
74 lines
2.4 KiB
Python
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))
|