"""
Experimental. Tools for visualizing the torch.jit.Graph objects.
"""
import string
import json
_vis_template = string.Template("""
$name
""")
def write(self, filename):
"""
Write an html file that visualizes a torch.jit.Graph using vis.js
Arguments:
self (torch.jit.Graph): the graph.
filename (string): the output filename, an html-file.
"""
nodes = []
edges = []
options = {}
for n, i in enumerate(self.inputs()):
nodes.append({
'id': i.unique(),
'label': 'input {}'.format(n),
'shape': 'square',
})
existing = set()
def add_edge(i_, n):
i = i_ if i_.kind() != 'Select' else i_.input()
if (i, n) in existing:
return
existing.add((i, n))
e = {
'from': n.unique(),
'to': i.unique(),
'arrows': 'from',
}
edges.append(e)
counts = {}
offset = 0
for n in self.nodes():
if len(n.uses()) == 0 or n.kind() == 'Undefined':
continue
ident = counts.get(n.kind(), 0)
counts[n.kind()] = ident + 1
d = {
'id': n.unique(),
'label': '{}_{}'.format(n.kind(), ident),
'y': offset,
'fixed': {'y': True},
}
if n in self.outputs():
d['shape'] = 'triangle'
for i in n.inputs():
add_edge(i, n)
nodes.append(d)
offset += 30
result = _vis_template.substitute(nodes=json.dumps(nodes),
edges=json.dumps(edges),
options=json.dumps(options),
name=filename)
with open(filename, 'w') as f:
f.write(result)