Skip to content

Instantly share code, notes, and snippets.

@apaszke
Forked from szagoruyko/pytorch-graphviz.py
Last active October 19, 2020 05:04
Show Gist options
  • Save apaszke/01aae7a0494c55af6242f06fad1f8b70 to your computer and use it in GitHub Desktop.
Save apaszke/01aae7a0494c55af6242f06fad1f8b70 to your computer and use it in GitHub Desktop.
from graphviz import Digraph
dot = Digraph(comment='LRP', node_attr={'style': 'filled', 'shape': 'box'})#, 'fillcolor': 'lightblue'})
seen = set()
def add_nodes(var):
if var not in seen:
if isinstance(var, Variable):
dot.node(str(id(var)), str(var.size()), fillcolor='lightblue')
else:
dot.node(str(id(var)), type(var).__name__)
seen.add(var)
if hasattr(var, 'previous_functions'):
for u in var.previous_functions:
dot.edge(str(id(u[0])), str(id(var)))
add_nodes(u[0])
add_nodes(R.creator)
dot.save('/tmp/lrp.dot')
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment