目标
为Relay IR生成图片形式的计算图。
实现方式
使用RelayVisualizer可视化Relay,RelayVisualizer定义了一组接口(包括渲染器、解析器)将IRModule可视化为节点和边,并且提供了默认解析器和渲染器。
首先需要安装依赖:
pip install graphviz
然后运行以下代码:
from tvm.driver import tvmc
from tvm.contrib import relay_viz# 加载模型为TVMCModule
model = tvmc.load('lenet5.onnx')# 使用RelayVisualizer可视化
viz = relay_viz.RelayVisualizer(model.mod)
viz.render("lenet5")
运行可以得到文本文件lenet5:
@main([Var(input.1, ty=TensorType([1, 1, 28, 28], float32))])
`--Call |--nn.relu `--Call |--add |--Call | |--nn.dense | |--Call | | |--nn.relu | | `--Call | | |--add | | |--Call | | | |--nn.dense | | | |--Call | | | | |--nn.relu | | | | `--Call | | | | |--add | | | | |--Call | | | | | |--nn.dense | | | | | |--Call | | | | | | |--reshape | | | | | | `--Call | | | | | | |--nn.max_pool2d | | | | | | `--Call | | | | | | |--nn.relu | | | | | | `--Call | | | | | | |--nn.bias_add | | | | | | |--Call | | | | | | | |--nn.conv2d | | | | | | | |--Call | | | | | | | | |--nn.max_pool2d | | | | | | | | `--Call | | | | | | | | |--nn.relu | | | | | | | | `--Call | | | | | | | | |--nn.bias_add | | | | | | | | |--Call | | | | | | | | | |--nn.conv2d | | | | | | | | | |--Var(Input) name_hint: input.1, shape: (1, 1, 28, 28), dtype: float32| | | | | | | | | `--Const shape: (6, 1, 5, 5), dtype: float32| | | | | | | | `--Const shape: (6,), dtype: float32| | | | | | | `--Const shape: (16, 6, 5, 5), dtype: float32| | | | | | `--Const shape: (16,), dtype: float32| | | | | `--Const shape: (120, 256), dtype: float32| | | | `--Const shape: (120,), dtype: float32| | | `--Const shape: (84, 120), dtype: float32| | `--Const shape: (84,), dtype: float32| `--Const shape: (10, 84), dtype: float32`--Const shape: (10,), dtype: float32
因为RelayVisualizer默认渲染器plotter = TermPlotter(),默认解析器parser = TermVizParser(),如果要生成图片格式的计算图,只需要设置RelayVisualizer的plotter和parser参数为Dot渲染器和Dot解析器:
viz = relay_viz.RelayVisualizer(model.mod,plotter=relay_viz.DotPlotter(),parser=relay_viz.DotVizParser())
运行会得到pdf文件lenet.pdf,截图如下:
可以为Relay自定义渲染器和解析器,有需要可以点击前往官网查看,在这里就不详细介绍了。
解释
DOT语言
DOT语言是一种文本图形描述语言,它提供了一种简单的描述图形的方法,并且可以为人类和计算机程序所理解。
Example.dot:
graph Example {A -- B;B -- C;C -- A;D -- A;
}
Graphviz
Graphviz(Graph Visualization Software)是一个由AT&T实验室启动的开源工具包,用于绘制DOT语言脚本描述的图形。
Graphviz运行:
dot -Tpng Example.dot -o Example.png
可以将Example.dot绘制为png格式的图片: