Explorador de modelos: simplificación de los modelos de AA para dispositivos perimetrales

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

Nos complace compartir el Explorador de modelos, una poderosa herramienta de visualización de gráficos. Ya sea que estés desarrollando tu propio modelo, convirtiendo un modelo entre formatos, optimizando para un dispositivo específico o depurando el rendimiento y la calidad, la capacidad de visualizar la arquitectura del modelo y cómo fluyen los datos entre los nodos puede ser muy útil. Con una visualización intuitiva y jerárquica de hasta los gráficos más grandes, el Explorador de modelos permite a los desarrolladores superar las complejidades de trabajar con modelos grandes, especialmente cuando se optimizan para dispositivos perimetrales.

Esta es la tercera entrada de blog de nuestra serie en la que presentamos los lanzamientos de desarrolladores de Google AI Edge: en las dos primeras, presentamos AI Edge Torch y la API generativa, que permiten usar los modelos de PyTorch y los LLM de alto rendimiento integrados en el dispositivo.

El Explorador de modelos, desarrollado originalmente como una utilidad para investigadores e ingenieros de Google, ahora está disponible públicamente como parte de nuestra familia de productos Google AI Edge. La versión inicial del Explorador de modelos ofrece lo siguiente:

  • Visualización de gráficos de modelos muy grandes, que contienen decenas de miles de nodos, en una experiencia fluida de 60 FPS

  • Ejecución directa en las notebooks de Colab

  • Extensibilidad a más formatos de modelos a través de extensiones personalizadas

  • Superposiciones de metadatos (por ejemplo, atributos, entradas/salidas, etc.) y datos personalizados (por ejemplo, rendimiento) directamente en los nodos

  • Potentes funciones de IU, incluida la vista de panel dividido para la comparación de gráficos, diseñadas para ayudarte a trabajar más rápido

En esta entrada de blog, explicaremos cómo comenzar a usar el Explorador de modelos y veremos un par de casos de uso útiles. Puedes encontrar más documentación y ejemplos aquí.


Cómo comenzar

El paquete PyPI del Explorador de modelos, de instalación sencilla, se ejecuta localmente en tu dispositivo, en Colab o en un archivo Python.


Ejecución local en tu dispositivo

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

Con estos comandos, se inicia un servidor en localhost:8080 y se abre la app web del Explorador de modelos en una pestaña del navegador. Obtén más información sobre el uso de las líneas de comandos del Explorador de modelos en la guía de línea de comandos.

Cuando tengas el servidor localhost en ejecución, carga tu archivo de modelo desde tu computadora (los formatos compatibles incluyen los utilizados por JAX, PyTorch, TensorFlow y TensorFlow Lite) y, si es necesario, selecciona el mejor adaptador para tu modelo en el menú desplegable "Adaptador". Haz clic aquí para obtener información sobre cómo utilizar el sistema de extensión del adaptador del Explorador de modelos para visualizar formatos de modelo no compatibles.


Ejecutar en una notebook de Colab

# Download a model (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("https://storage.googleapis.com/tfweb/model-graph-vis-v2-test-models/efficientdet.tflite", model_path)
 
# Install Model Explorer
pip install ai-edge-model-explorer
 
# Visualize the downloaded EfficientDet model
import model_explorer
model_explorer.visualize(model_path)

Después de ejecutar la celda, el Explorador de modelos se muestra en un iFrame en una nueva celda. En Chrome, también aparece el botón "Abrir en una nueva pestaña". Si haces clic en él, puedes mostrar la IU en una pestaña separada. Consulta aquí para obtener más información sobre cómo ejecutar el Explorador de modelos en Colab.

Visualize models via the Model Explorer API

Abre el visualizador con las APIs del Explorador de modelos

El paquete model_explorer proporciona APIs convenientes que te permiten visualizar modelos de archivos o de un módulo PyTorch, y una API de nivel inferior para visualizar modelos de múltiples fuentes. Asegúrate de instalarlo primero siguiendo la guía de instalación. Consulta la guía de la API del Explorador de modelos para obtener más información.

A continuación, se muestra un ejemplo de cómo visualizar un modelo de PyTorch. La visualización de los modelos de PyTorch requiere un enfoque ligeramente diferente a otros formatos porque PyTorch no tiene un formato de serialización estándar. El Explorador de modelos ofrece una API especializada para visualizar los modelos de PyTorch de forma directa, utilizando el ExportedProgram de torch.export.export.

import model_explorer
import torch
import torchvision
 
# Prepara un modelo PyTorch y sus entradas
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 importa de qué manera visualices tus modelos, el Explorador de modelos implementa en segundo plano la representación de gráficos acelerada por GPU con WebGL y three.js a fin de lograr una experiencia de visualización fluida de 60 FPS, incluso con gráficos que contienen decenas de miles de nodos. Si te interesa obtener más información sobre cómo el Explorador de modelos renderiza gráficos grandes, lee el blog de Google Research.


Comprender la arquitectura del modelo

Los modelos grandes pueden ser complejos, pero el Explorador de modelos los hace más fáciles de entender, ya que divide la visualización en capas jerárquicas. Echa un vistazo al modelo MobileBert, que se muestra a continuación: está claro cómo la máscara de autoatención y la incorporación se introducen en una capa de transformador. Incluso puedes profundizar en la capa de incorporación para comprender las relaciones entre los diferentes tipos de incorporaciones. La vista jerárquica del Explorador de modelos hace que hasta las arquitecturas de modelos más complejas sean más fáciles de entender.

Model architecture
Con información de capas, como `self_attention_mask`, `embedding` y `transformer_layer`, es fácil entender la arquitectura del modelo.

Depurar errores de conversión

Convertir modelos de un formato a otro (como de PyTorch a TFLite) puede ser complicado, pero el Explorador de modelos te ayuda a comparar los gráficos originales y convertidos. De esta manera, es fácil detectar cualquier cambio que pueda afectar el rendimiento de tu modelo. Por ejemplo, en las imágenes que se muestran a continuación, puedes ver cómo cambió un subgráfico dentro de una capa durante la conversión, lo que te ayuda a identificar y corregir posibles errores.

A side-by-side comparison of the original PyTorch model (left) and the optimized TensorFlow Lite model (right)
Una comparación del modelo original de PyTorch (izquierda) y el modelo optimizado de TensorFlow Lite (derecha), útil para depurar la conversión de gráficos del modelo.

Depurar el rendimiento y la precisión numérica con superposiciones de datos por nodo

Una característica clave del Explorador de modelos es su capacidad para superponer datos por nodo en un gráfico, lo que permite ordenar, buscar y personalizar nodos con los valores de esos datos. Combinada con el sistema de vista jerárquica del Explorador de modelos, esta función te permite detectar rápidamente el rendimiento o los cuellos de botella numéricos. En el siguiente ejemplo, se muestra el error cuadrático medio en cada nodo entre un modelo TFLite cuantificado y su contraparte de punto flotante. El Explorador de modelos destaca que la caída de calidad está cerca de la parte inferior del gráfico, lo que te brinda la información que necesitas para ajustar tu método de cuantificación. Para obtener más información sobre cómo trabajar con datos personalizados en el Explorador de modelos, consulta nuestra documentación detallada en Github.

Per-node data overlay
Esta superposición de datos por nodo permite a los usuarios identificar rápidamente problemas numéricos o de rendimiento dentro de un modelo.

Uso en Google

Entre los usuarios más destacados del Explorador de modelos en Google, se encuentran Waymo y Google Silicon. El Explorador de modelos desempeñó un papel crucial para ayudar a estos equipos a depurar y optimizar modelos integrados en el dispositivo, como 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

Lo que viene

Este es solo el comienzo. En los próximos meses, nos centraremos en mejorar la parte más importante: refinaremos las funciones clave de la IU, como la diferenciación y edición de gráficos, y mejoraremos la extensibilidad permitiéndote integrar tus propias herramientas en el Explorador de modelos.



Agradecimientos

Este trabajo es una colaboración entre múltiples equipos funcionales de Google. Nos gustaría extender nuestro agradecimiento a los ingenieros 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, a nuestra gerente del programa técnico, Kristen Wright, y a nuestro gerente de producto, Aaron Karp. También queremos agradecer al equipo de UX, que incluye a Zi Yuan, Anila Alexander, Elaine Thai, Joe Moran y Amber Heinbockel.