计图团队联合发布可视化工具JittorVis,帮助开发者理解深度神经网络模型

6月30日,清华大学计图(Jittor)团队与刘世霞教授课题组合作开发的基于Jittor的可视化工具JittorVis正式发布。该工具用于帮助理解基于Jittor构建的深度神经网络模型的计算图结构和数据流信息。

Part 1

计算图可视化引起学界关注

近年来,深度神经网络在图像识别、检测、分割、生成等方面取得了突破性的进展。而随着数据量的增加,网络模型的结构也越来越复杂。一个大规模神经网络可能会包含上百个网络中间层,每层可能包含上百万个参数。

深度神经网络结构的复杂性使得开发者们难以理解模型的内部结构。他们在使用模型的时候往往将其当作一个黑盒子,对模型内部层次的作用缺乏一个清晰的认知,通常需要对模型反复试验来得到一个较好的模型。

可视化是帮助开发者理解深度神经网络的有效手段,如发表于IEEE VIS 2016的CNNVis[1]首次提出了深度卷积神经网络可视化技术,显著提高了深度学习模型的可解释性,是ESI高被引论文,WOS引用在自2016年以来可视化领域发表的4195篇期刊论文中排名最高,Google Research发表于IEEE VIS 2017的TensorFlow Graph Visualizer[2]荣获最佳论文奖。

主流的深度学习框架均有对应的可视化工具,例如TensorFlow的TensorBoard[2](图1)。

图1 Google推出的计算图可视化工具TensorBoard

Part 2

JittorVis支持对模型更深入的理解

Jittor[3]是一个采用元算子表达神经网络计算单元、完全基于动态编译的深度学习框架,由于其性能上的优越性而广受关注。

目前已有的计算图可视化工具不能深入到元算子级别的分析,无法直接用于Jittor,导致Jittor开发者无法通过已有的可视化方法有效地开发深度神经网络模型。

为此,JittorVis集成了CNNVis[1]和DGMTracker[4]的优势,直观地展示深度神经网络模型的计算图结构和数据流信息,为开发者们提供模型内部的层次结构信息,让开发者可以从整体到局部对模型有深入的理解,方便开发者更有效地调试和修改模型,不断提高模型性能。

JittorVis的主要特性包括:

  • 展示模型内部的层次结构信息

  • 在不同级别(从元算子到计算节点)上对计算图进行可视化

  • 剖析Jittor程序

图2 JittorVis可视化界面

JittorVis包含三个视图(见图2):

  • 统计信息视图

位于界面上方,展示模型训练过程中的统计信息,例如损失函数和准确率曲线。横轴表示训练轮次,纵轴表示评价指标。

  • 层次结构视图

位于界面的左侧,展示模型内部的层次结构信息。每个节点表示计算图中的一个计算节点,每个叶节点表示计算图中的一个元算子。

  • 图结构视图

位于界面的右侧,展示模型的计算图。每个矩形表示一个计算节点,矩形之间的连接表示计算节点之间的数据流。

Part 3

JittorVis支持对模型的交互式探索

为了帮助开发者更好地理解模型的内部工作机理,JittorVis主要支持以下交互:

1)展示不同评价指标在训练过程中的变化

在统计信息视图中,浏览训练过程中模型的不同评价指标变化有助于开发者全面理解模型性能变化。

已支持的评价指标包括损失(loss),准确率(accuracy),查准率(precision),查全率(recall),内存占用(memory)。

2)选择并展示一部分计算节点

一个深度神经网络模型可能包含数以千计的计算节点。因此,支持开发者选取一部分计算节点进行局部分析是很有必要的。层次结构视图允许用户选取一部分计算节点进行详细分析(图3)。

相应地,图结构视图展示所选择计算节点的内部结构、特征图以及计算节点之间的数据流走向(图4)。

这两个视图协同工作,帮助用户更好地理解Jittor模型的工作机理,对模型的训练和决策过程进行诊断,进而实现模型的改进。

图3 ResNet模型中的层次结构

图4 ResNet“layer2”的计算图算子结构

目前,Jittor的可视化工具JittorVis已经开源,开源链接为:

欢迎大家使用Jittor的可视化工具JittorVis开展研究工作。

如果大家在使用过程中发现有什么问题,请大家在github上提交issue或者pull request。也欢迎加入Jittor开发者的QQ交流群,随时提问和交流!期待您提出宝贵的意见。

参考文献

  1. Liu M., Shi J., Li Z., Li C., Zhu J., & Liu S. Towards better analysis of deep convolutional neural networks, IEEE Transactions on Visualization and Computer Graphics, 23(1), 91-100, 2017.
  2. Wongsuphasawat K., Smilkov D., Wexler J., Wilson J., Mane D., Fritz D., Krishnan D., Viegas F. B. and Wattenberg M., Visualizing dataflow graphs of deep learning models in tensorflow, IEEE Transactions on Visualization and Computer Graphics, 24(1), 1–12, 2018.
  3. Hu S.-M., Liang D., Yang G.-Y., Yang G.-W., & Zhou W.-Y. Jittor: a novel deep learning framework with meta-operators and unified graph execution, Science China Information Sciences, 63(12), 1-21, 2020.
  4. Liu M., Shi J., Cao K., Zhu J. & Liu S, Analyzing the training processes of deep generative models, IEEE Transactions on Visualization and Computer Graphics, 24(1), 77-87, 2018.