pytorch-graphviz.py

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)), str(type(var)))
        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')