Model Explorer:针对边缘设备简化 ML 模型

六月 26, 2024
Kristen Wright Technical Program Manager
Eric Yang Software Engineer

我们很高兴能够分享一款功能强大的图形可视化工具:Model Explorer。无论您是开发自己的模型、转换模型格式、针对特定设备进行优化,还是调试性能和质量,借助 Model Explorer 直观了解模型架构以及数据在节点之间如何流动都会对您大有帮助。通过以直观、分层的方式呈现甚至超大的图形,Model Explorer 让开发者能够克服处理大模型的复杂性,特别是针对边缘设备进行优化时。

这是我们系列博文中关于 Google AI Edge 开发者版本的第三篇:前两篇博文介绍了 AI Edge Torch生成式 API,用于在设备上启用 PyTorch 模型和高性能 LLM。

Model Explorer 最初是面向 Google 研究人员和工程师开发的实用程序,现在已作为 Google AI Edge 系列产品的一部分公开提供。Model Explorer 的初始版本提供以下内容:

  • 即使包含数万个节点的超大模型图形,也能实现流畅的 60 FPS 可视化体验

  • 直接在 Colab notebook 中执行

  • 借助自定义扩展功能,扩展至更多模型格式

  • 直接在节点上覆盖元数据(例如,属性、输入/输出等)和自定义数据(例如,性能)

  • 旨在帮助您提高工作速度的强大界面功能,包括用于并排图形比较的拆分窗格视图

在本博文中,我们将介绍如何开始使用 Model Explorer 和几个强大的用例。请点击此处查看更多文档和示例。


开始构建

Model Explorer 易于安装的 PyPI 软件包可在您的设备、Colab 或 Python 文件中本地运行。


在您的设备上本地运行

$ pip install ai-edge-model-explorer
$ model-explorer
 
 http://localhost:8080 中启动 Model Explorer 服务器

这些命令将在 localhost:8080 中启动服务器,并在浏览器标签页中打开 Model Explorer 网络应用。有关 Model Explorer 命令行用法的更多信息,请阅读命令行指南

运行本地主机服务器后,利用计算机上传模型文件(支持的格式包括 JAXPyTorchTensorFlowTensorFlow Lite 使用的格式)。如有需要,请通过“适配器”下拉菜单为您的模型选择最佳适配器。请点击此处,了解如何利用 Model Explorer 适配器扩展系统来可视化不支持的模型格式。


在 Colab notebook 中运行

# 下载模型(此示例中使用的是 Efficientdet TFLite 模型)
import os
import tempfile
import urllib.request
 
tmp_path = tempfile.mkdtemp()
model_path = os.path.join(tmp_path, 'model.tflite')
urllib.request.urlretrieve("https://storage.googleapis.com/tfweb/model-graph-vis-v2-test-models/efficientdet.tflite", model_path)
 
# 安装 Model Explorer
pip install ai-edge-model-explorer
 
# 可视化下载的 EfficientDet 模型
import model_explorer
model_explorer.visualize(model_path)

运行单元后,Model Explorer 将在嵌入新单元的 iFrame 中显示。在 Chrome 中,界面还会显示“在新标签页中打开”按钮,您点击该按钮即可让界面在单独的标签页中显示。请访问此处,了解有关在 Colab 中运行 Model Explorer 的更多信息。

Visualize models via the Model Explorer API

使用 Model Explorer API 打开可视化工具

model_explorer 软件包提供使用方便的 API,可让您可视化来自文件或 PyTorch 模块的模型;您也可以利用其中更低级的 API 来可视化来自多个来源的模型。请务必先按照安装指南安装 API。请查看 Model Explorer API 指南,以了解更多信息。

以下是如何可视化 PyTorch 模型的示例。PyTorch 没有标准的序列化格式,因为可视化 PyTorch 模型需要使用与其他格式略微不同的方法。Model Explorer 提供了一个专门的 API,可使用来自 torch.export.exportExportedProgram 直接可视化 PyTorch 模型。

import model_explorer
import torch
import torchvision
 
# 准备 PyTorch 模型及其输入 
model = torchvision.models.mobilenet_v2().eval()
inputs = (torch.rand([1, 3, 224, 224]),)
ep = torch.export.export(model, inputs)
 
# 可视化
model_explorer.visualize_pytorch('mobilenet', exported_program=ep)

无论您以哪种方式可视化模型,Model Explorer 都能在后台使用 WebGL 和 three.js 实现 GPU 加速的图形渲染,即使图形包含数万个节点,也能实现流畅的 60 FPS 可视化体验。如果您有兴趣了解更多关于 Model Explorer 如何渲染大图形的信息,您可以阅读 Google Research 博客中的相关内容。


了解模型架构

大模型可能很复杂,但 Model Explorer 通过以分层方式分解模型的可视化效果,使它们更容易理解。请看下图所示的 MobileBert 模型:自注意掩码和嵌入是如何馈送到 Transformer 层的一目了然。您甚至可以更深入地研究嵌入层,以理解不同类型的嵌入之间的关系。Model Explorer 的分层视图让即使是最复杂的模型架构也更容易理解。

Model architecture
使用层信息,如“self_attention_mask”、“embedding”和“transformer_layer”,可以轻松理解模型架构。

调试转换错误

将模型从一种格式转换为另一种格式(例如,从 PyTorch 转换为 TFLite)可能很棘手,但 Model Explorer 可帮助您并排比较原始图形和转换后的图形。这样,您便可以轻松发现可能影响模型性能的任何更改。例如,在下面的图像中,您可以看到图层中的子图在转换过程中是如何变化的,从而帮助您识别和修复潜在的错误。

A side-by-side comparison of the original PyTorch model (left) and the optimized TensorFlow Lite model (right)
原始 PyTorch 模型(左)和经过优化的 TensorFlow Lite 模型(右)的并排比较 – 这在调试模型图形转换时十分有用。

使用每节点数据覆盖来调试性能和数字准确度

Model Explorer 有一项强大的功能,它能够覆盖图形上的每节点数据,让您可以使用该数据中的值对节点进行排序、搜索和设置样式。与 Model Explorer 的分层视图系统相结合,此功能使您能够快速发现性能或数字瓶颈。下面的示例显示了在经量化的 TFLite 模型与其浮点对应模型之间,每个节点的均方误差。Model Explorer 会突出显示质量下降是否接近图形底部的位置,从而为您提供调整量化方法所需的信息。要了解有关在 Model Explorer 中使用自定义数据的更多信息,请查看我们在 Github 上的详细介绍文档。

Per-node data overlay
此每节点数据覆盖允许用户快速确定模型中的性能或数字问题。

在 Google 的使用情况

在 Google,最常使用 Model Explorer 的用户是 Waymo 和 Google Silicon。Model Explorer 在帮助这些团队调试和优化 Gemini Nano 等设备端模型方面发挥了至关重要的作用。

Model Explorer...allows us to swiftly address bottlenecks, leading to the successful launch of multiple image, speech, and LLM use cases - especially the Gemini Nano on Pixel devices" - Dongdong Li - Pixel AI Experiences Lead

未来计划

我们认为这只是一个开始。在接下来的几个月里,我们将专注于通过优化图形区分和编辑等关键界面功能来增强核心,并通过允许您将自己的工具集成到 Model Explorer 中来提高可扩展性。



致谢

这项工作由 Google 的多个职能团队之间合作完成。在此,我们要感谢工程师 Na Li、Jing Jin、Eric (Yijie) Yang、Akshat Sharma、Chi Zeng、Jacques Pienaar、Chun-nien Chan、Jun Jiang、Matthew Soulanille、Arian Arfaian、Majid Dadashi、Renjie Wu、Zichuan Wei、Advait Jain、Ram Iyengar、Matthias Grundmann、Cormac Brick、Ruofei Du、技术项目经理 Kristen Wright 以及产品经理 Aaron Karp。我们还要感谢用户体验团队的成员,包括 Zi Yuan、Anila Alexander、Elaine Thai、Joe Moran 和 Amber Heinbockel。