目标

为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格式的图片：