图形可视化在机器学习（ML）开发过程中起着关键作用。模型的可视化表示有助于研究人员和工程师调试转换和量化问题，识别性能瓶颈，找到优化模式（例如，操作融合），并深入了解模型架构。这些做法对许多类型的模型都很有用，尤其是在将模型部署到资源有限的设备（如手机和浏览器）时。然而，现代ML和AI模型（如变压器和扩散器）的规模和复杂性不断增加，对现有的可视化工具提出了重大挑战。基于Transformer的大规模模型往往会淹没传统的图形可视化工具，导致渲染失败或视觉复杂性难以管理。此外，缺乏清晰的视觉层次使人们对节点关系的理解变得复杂。
为了解决这些限制，我们引入了Model Explorer，这是一种新颖的图形可视化解决方案，可以平滑地处理大型模型，并可视化层次信息，如函数名称和作用域。Model Explorer支持多种图形格式，包括JAX、PyTorch、TensorFlow和TensorFlow Lite使用的图形格式。Model Explorer最初是为谷歌研究人员和工程师开发的实用程序，现在作为我们谷歌AI Edge系列产品的一部分公开提供。
Model Explorer证明其在精简大型模型到设备平台的部署方面特别有效，在设备平台上，可视化转换、量化和优化数据尤其有用。我们描述了模型资源管理器如何将3D游戏和动画制作中使用的图形技术（如实例化渲染和多通道有符号距离域（MSDF））相结合，并将其用于ML图形渲染。
大型图形可视化面临的挑战
可视化大型模型图有两个主要的技术挑战。首先，布局算法很难与更大的图进行缩放：随着节点数量的增长，其计算复杂度急剧上升，导致布局阶段的速度显著减慢，偶尔会出现完全故障。即使是成功的布局也往往过于密集和复杂，降低了其可解释性。其次，大多数现有的模型可视化工具使用基于可缩放矢量图形的渲染，这不是为渲染大量对象而设计的。缩放和滚动操作变得迟缓且无响应，导致可视化工具无法使用。
受Tensorboard图形可视化工具中分层布局概念的启发，我们开发了一个库，可以从主流创作框架（如TensorFlow、PyTorch和JAX）中提取分层信息。然后，我们构建了一个交互式系统，从最顶层开始可视化节点，允许用户通过逐个展开层来逐步导航图形，或者通过折叠层来恢复到更高级别的层。这允许用户以所需的粒度检查图中的内部结构和连接。由于布局算法是按层操作的，因此当用户选择打开层时，模型资源管理器会计算布局计算。该方法避免了在初始图形加载过程中对所有节点进行不必要的计算，并显著提高了大型模型的性能。
我们通过使用WebGL和three.js实现GPU加速的图形渲染来解决第二个挑战。因此，我们实现了每秒60帧（FPS）的平滑用户体验，即每秒在屏幕上显示60个单独的图像，在交互和动画中创建流畅而逼真的运动，即使图形包含数万个节点。此外，我们还利用实例化渲染技术在场景中的不同位置同时渲染对象的多个副本。下面的示例显示了具有50000个节点和5000条边（出于演示目的随机生成）的图形是如何在带有集成GPU的2019 Macbook Pro上以60 FPS的速度平滑渲染的。为了提高视觉连续性，我们为图层导航添加了平滑的动画，帮助用户保持对模型结构中位置的理解。
了解模型体系结构
使用基于图层的视图和导航复杂结构的能力，大型模型变得更容易理解。例如，在下面的MobileBert模型中，很明显，自注意掩码和嵌入被馈送到转换器层中。嵌入层的展开视图显示了不同类型嵌入之间的关系。该模型总共有近2000个节点，如果没有这些层次信息，几乎不可能理解。
调试转换错误
在部署到特定硬件（如手机或笔记本电脑）之前，ML模型必须经过转换过程：例如，将PyTorch模型转换为Tensorflow Lite模型。但是，在转换过程中，有关转换的信息往往会丢失。为了帮助比较多个图形，模型资源管理器提供了并排比较模式。下面的示例突出显示了一个层内的子图（即torc.nn.modules.linear_attn）在从PyTorch到TensorFlow Lite的转换过程中是如何变化的。比较层的输入和输出的形状和数据类型信息可以帮助突出显示转换错误。
调试性能和数字精度
模型资源管理器的另一个功能是它能够在图上覆盖每个节点的数据，允许用户使用数据中的值对节点进行排序、搜索和样式化。结合分层视图，我们的用户研究参与者分享说，这有助于他们快速将性能或数字问题缩小到图表中的某个区域，如果没有这种可视化，这些见解将更加难以捉摸。以下示例显示了量化TFLite模型相对于其浮点对应模型的均方误差。使用模型浏览器，用户能够快速识别质量下降接近图表底部，并调整量化方法。
结论
模型资源管理器引入了一种强大的新方法来检查几乎任何大小的模型中的体系结构和调试问题，而不会影响用户体验或渲染性能。它清晰地呈现了模型结构，使用层和分组来提高理解，并结合了调试功能和层级见解来支持模型分析。
说明
模型资源管理器现已公开。有关安装说明和提供反馈，请访问模型资源管理器网站。
