TVM Relay IR计算图可视化

from.expr_functorimportExprFunctor
from.importexpras_expr
importnetworkxasnx
classVisualizeExpr(ExprFunctor):
def__init__(self):
super().__init__()
self.graph=nx.DiGraph()
self.counter=0
defviz(self,expr):
assertisinstance(expr,_expr.Function)
forparaminexpr.params:
self.visit(param)
returnself.visit(expr.body)
defvisit_constant(self,const):# overload this!
pass
defvisit_var(self,var):
name=var.name_hint
self.graph.add_node(name)
self.graph.nodes[name]['style']='filled'
self.graph.nodes[name]['fillcolor']='mistyrose'
returnvar.name_hint
defvisit_tuple_getitem(self,get_item):
tuple=self.visit(get_item.tuple_value)
# self.graph.nodes[tuple]
index=get_item.index
# import pdb; pdb.set_trace()
returntuple
defvisit_call(self,call):
parents=[]
forargincall.args:
parents.append(self.visit(arg))
# assert isinstance(call.op, _expr.Op)
name="{}({})".format(call.op.name,self.counter)
self.counter+=1
self.graph.add_node(name)
self.graph.nodes[name]['style']='filled'
self.graph.nodes[name]['fillcolor']='turquoise'
self.graph.nodes[name]['shape']='diamond'
edges=[]
fori,parentinenumerate(parents):
edges.append((parent,name,{'label':'arg{}'.format(i)}))
self.graph.add_edges_from(edges)
returnname
defvisualize(expr,mydir="relay_ir.png"):
viz_expr=VisualizeExpr()
viz_expr.viz(expr)
graph=viz_expr.graph
dotg=nx.nx_pydot.to_pydot(graph)
dotg.write_png(mydir)

fromtvm.relay.visualizeimportvisualize
func=relay.Function(relay.analysis.free_vars(output),output)
visualize(func)