目标

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