当前位置:首页|资讯

【TVM 教程】使用 Relay Visualizer 可视化 Relay

作者:HyperAI超神经发布时间:2024-09-26

Apache TVM 是一个端到端的深度学习编译框架,适用于 CPU、GPU 和各种机器学习加速芯片。更多 TVM 中文文档可访问 →https://tvm.hyper.ai/

作者:Chi-Wei Wang

Relay IR 模块可以包含很多操作。通常单个操作很容易理解,但放在一起可能会使计算图难以阅读。随着优化 pass 发挥作用,情况可能会变得更糟。

该实用程序定义了一组接口(包括解析器、绘图器/渲染器、计算图、节点和边)将 IR 模块可视化为节点和边,且提供了默认解析器。用户可以用自己的渲染器对计算图进行渲染。

这里用渲染器来渲染文本形式的计算图,它是一个轻量级、类似 AST 可视化工具(灵感来自 clang ast-dump)。以下将介绍如何通过接口类来实现自定义的解析器和渲染器。

安装依赖可运行:

更多细节参考 tvm.contrib.relay_viz


定义具有多个 GlobalVar 的 Relay IR 模块

构建一个包含多个 GlobalVar 的 Relay IR 模块示例。定义一个 add 函数,并在 main 函数中调用。


在终端上使用 Relay Visualizer 渲染图形

终端可以用类似于 clang AST-dump 的文本显示 Relay IR 模块,可以看到 AddFunc 函数在 main 函数中调用了两次。

输出结果:

为 Relay 类型自定义解析器

有时想要强调感兴趣的信息,或者针对特定用途对事物进行不同的解析,需要遵守接口来定制解析器,下面演示如何自定义 relay.var 的解析器,需要实现抽象接 口 tvm.contrib.relay_viz.interface.VizParser

将解析器和感兴趣的渲染器传递给 visualizer,这里只用终端渲染器。

输出结果:


围绕计算图和绘图器进行定制

除了解析器,还可以通过实现抽象类 tvm.contrib.relay_viz.interface.VizGraph 和 tvm.contrib.relay_viz.interface.Plotter 来自定义计算图和渲染器。下面重写了 terminal.py 中定义的 TermGraph。我们添加了一个 hook 在 AwesomeVar 上方复制,并让 TermPlotter 使用新类。

输出结果:


总结

本教程演示了 Relay Visualizer 的使用和自定义。 tvm.contrib.relay_viz.RelayVisualizer 类由 interface.py 中定义的接口组成。

目的是快速查看,然后修复迭代,构造函数参数尽可能简单,而自定义可以通过一组接口类进行。

下载 Python 源代码:using_relay_viz.py

下载 Jupyter Notebook:using_relay_viz.ipynb





Copyright © 2024 aigcdaily.cn  北京智识时代科技有限公司  版权所有  京ICP备2023006237号-1