Model Explorer: Simplifying Ml Models For Edge Devices

Sedang Trending 1 bulan yang lalu

We are excited to share Model Explorer - a powerful graph visualization tool designed to help you understand and debug your ML models. With an intuitive, hierarchical visualization of even the largest graphs, Model Explorer enables developers to overcome the complexities of optimizing models for edge devices. This is the third blog post in our series covering Google AI Edge developer releases. If you missed the first two, be sure to check out the AI Edge Torch and Generative API blogs.

Developed originally as a utility for Google researchers and engineers, Model Explorer is now publicly available as part of our Google AI Edge family of products. The initial version of Model Explorer offers the following:

  • GPU-based rendering engine to visualize large model graphs
  • Popular ML framework support
  • Runs directly in Colab notebooks
  • Adapter extension system to visualize more model formats
  • Overlay metadata (e.g., attributes, inputs/outputs, etc) and custom-data (e.g. performance) directly on nodes
  • Powerful UI feature suite designed to help you work faster

In this blog post we will walk through how to get started with Model Explorer and how to utilize Model Explorer’s custom info overlay API to debug and optimize your models. Further documentation and examples are available here.

Getting started

Model Explorer prioritizes a seamless user experience. Its easy-to-install PyPI package runs locally on your device, in Colab, and in a Python file, boosting the privacy and security of your model graphs.

Run locally on your device

$ pip install ai-edge-model-explorer $ model-explorer Starting Model Explorer server at http://localhost:8080

These commands will start a server at localhost:8080 and open the Model Explorer web app in a browser tab. See more info about Model Explorer command lines in the command line guide.

Once you have a localhost server running, upload your model file from your computer (formats supported include those used by JAX, PyTorch, TensorFlow and TensorFlow Lite) and select the best adapter for your model via the ‘Adapter’ drop down menu on the home page. Visit here to learn how to utilize the Model Explorer adapter extension system to visualize unsupported model formats.

Run in Colab notebooks

# Download your models (this example uses an Efficientdet TFLite model) import os import tempfile import urllib.request tmp_path = tempfile.mkdtemp() model_path = os.path.join(tmp_path, 'model.tflite') urllib.request.urlretrieve("", model_path) # Install Model Explorer pip install ai-edge-model-explorer # Visualize the downloaded EfficientDet model import model_explorer model_explorer.visualize(model_path)

After running the cell, Model Explorer will be displayed in an iFrame embedded in a new cell. In Chrome, the UI will also show an "Open in new tab" button that you can click to show the UI in a separate tab. Visit here to learn more about running Model Explorer Colab.

Visualize models via the Model Explorer API

Visualize models via the Model Explorer API

The model_explorer package provides convenient APIs to let you visualize models from files or from a PyTorch module, and a lower level API to visualize models from multiple sources. Make sure to install it first by following the installation guide. To learn more check out the Model Explorer API guide.

Below is an example for visualizing PyTorch models. Visualizing PyTorch models requires a slightly different approach due to their lack of a standard serialization format. Model Explorer offers a specialized API to visualize PyTorch models directly, using the ExportedProgram from torch.export.export.

import model_explorer import torch import torchvision # Prepare a PyTorch model and its inputs model = torchvision.models.mobilenet_v2().eval() inputs = (torch.rand([1, 3, 224, 224]),) ep = torch.export.export(model, inputs) # Visualize model_explorer.visualize_pytorch('mobilenet', exported_program=ep)

No matter which way you visualize your models, under the hood Model Explorer implements GPU-accelerated graph rendering with WebGL and three.js that achieves a smooth, 60 FPS visualization experience even with graphs containing tens of thousands of nodes.

Debug performance and numeric accuracy with node info overlay

A key Model Explorer feature is its ability to overlay per-node info on a graph, allowing you to sort, search, and stylize nodes using the values in that data. Combined with the hierarchical view, per-node info overlay enables you to quickly narrow down performance or numeric bottlenecks. The example below shows the mean squared error of a quantized TFLite model versus its floating point counterpart. Using Model Explorer, you're able to quickly identify that the quality drop is near the bottom of the graph, and adjust your quantization method as needed. Let’s walk through how to prepare and visualize custom node data.

Per-node info overlay

This per-node info overlay allows users to quickly identify performance or numeric issues within a model.

Prepare custom node data

We provide a set of Python APIs to help you create custom node info and serialize it into a JSON file. From a high level, the custom node info has the following structure:

ModelNodeData: The top-level container storing all the info for a model. It consists of one or more GraphNodeData objects indexed by graph ids.

GraphNodeData: Holds the info for a specific graph within the model. It includes:

  • results: Stores the custom node values, indexed by either node ids or output tensor names.
  • thresholds or gradient: color configurations that associate each node value with a corresponding node background color or label color, enabling visual representation of the data.

Below is a minimal example of preparing custom node info using the node_data_builder API. For in-depth documentation for preparing custom node info visit in our Github repo.

from model_explorer import node_data_builder as ndb # Populate values for the main graph in a model. main_graph_results: dict[str, ndb.NodeDataResult] = {} main_graph_results['node_id1'] = ndb.NodeDataResult(value=100) main_graph_results['node_id2'] = ndb.NodeDataResult(value=200) main_graph_results['any/output/tensor/name/'] = ndb.NodeDataResult(value=300) # Create a gradient color mapping. # # The minimum value in `main_graph_results` maps to the color with stop=0. # The maximum value in `main_graph_results` maps to the color with stop=1. # Other values maps to a interpolated color in-between. gradient: list[ndb.GradientItem] = [ ndb.GradientItem(stop=0, bgColor='yellow'), ndb.GradientItem(stop=1, bgColor='red'), ] # Construct the info for the main graph. main_graph_data = ndb.GraphNodeData( results=main_graph_results, gradient=gradient) # Construct the info for the model. model_data = ndb.ModelNodeData(graphsData={'main': main_graph_data}) # You can save the info to a json file. model_data.save_to_file('path/to/file.json')

Add per-node button in the app bar

You can also visualize custom node info by creating a config object and passing it to the visualize_from_config API.

import model_explorer from model_explorer import node_data_builder as ndb # Create a `ModelNodeData` as shown in previous section. model_node_data = ... # Create a config. config = model_explorer.config() # Add model and custom node info to it. (config .add_model_from_path('/path/to/a/model') # Add node info from a json file. # A node info json file can be generated by calling `ModelNodeData.save_to_file` .add_node_data_from_path('/path/to/node_data.json') # Add node info from info class object .add_node_data('my data', model_node_data)) # Visualize model_explorer.visualize_from_config(config)

Early Adoption

In the past few months, we have worked closely with early adoption partners including Waymo and Google Silicon to improve our visualization tool. Notably, Model Explorer has played a crucial role in helping these teams debug and optimize on-device models like Gemini Nano currently deployed in production.

"Allows us to swiftly address bottlenecks, leading to the successful launch of multiple image, speech, and LLM use cases - especially the Gemini Nano in Pixel devices" Dongdong Li – Pixel AI Experiences Lead

What’s next?

In the coming months we will focus on enhancing the core by refining key UI features like graph diffing and editing, empowering extensibility by allowing you to seamlessly integrate your own tools into Model Explorer, and open-sourcing the Model Explorer front-end. This is the third and final post of the AI Edge 3-part blog series. To stay up to date on the latest AI Edge updates visit the AI Edge site.


This work is a collaboration across multiple functional teams at Google. We would like to extend our thanks to engineers 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, our Technical Program Manager, Kristen Wright, and our Product Manager, Aaron Karp. We would also like to thank the UX team including Zi Yuan, Anila Alexander, Elaine Thai, Joe Moran and Amber Heinbockel.

Sumber Google Developers Blog
Google Developers Blog