Model Explorer: modelos de ML simplificados para dispositivos de borda

JUN 26, 2024

Kristen Wright Technical Program Manager

Eric Yang Software Engineer

Temos o prazer de compartilhar o Model Explorer, uma avançada ferramenta de visualização de grafos. Seja para desenvolver um modelo próprio, converter um modelo entre formatos diferentes, fazer a otimização para um dispositivo específico ou depurar o desempenho e a qualidade, a capacidade de visualizar a arquitetura do modelo e como os dados fluem entre os nós pode ser muito útil. Com uma visualização intuitiva e hierárquica até mesmo dos maiores grafos, o Model Explorer permite que os desenvolvedores superem as complexidades do trabalho com modelos grandes, particularmente na otimização para dispositivos de borda

Esta é a terceira postagem do blog de nossa série que aborda os lançamentos para desenvolvedores do Google AI Edge: as duas primeiras postagens apresentaram o AI Edge Torch e a API generativa, que ativam os modelos do PyTorch e os LLMs de alto desempenho no dispositivo.

Desenvolvido originalmente como um utilitário para pesquisadores e engenheiros do Google, o Model Explorer agora está disponível publicamente como parte da família de produtos Google AI Edge. A versão inicial do Model Explorer oferece:

Visualização de grafos de modelos muito grandes, contendo dezenas de milhares de nós, em uma experiência fluida de 60 QPS.

Suporte para frameworks populares de ML, incluindo JAX, PyTorch, TensorFlow e TensorFlow Lite.

Execução direta em notebooks do Colab.

Capacidade de extensão para mais formatos de modelo por meio de extensões personalizadas.

Sobreposições de metadados (por exemplo, atributos, entradas/saídas etc.) e dados personalizados (por exemplo, desempenho) diretamente em nós.

Recursos avançados de IU, incluindo visualização em painel dividido para comparação de grafos lado a lado, projetados para ajudar a acelerar seu trabalho.

Nesta postagem do blog, veremos como dar os primeiros passos com o Model Explorer e alguns casos de uso avançados. Mais documentação e exemplos estão disponíveis aqui.

Primeiros passos

O pacote PyPI fácil de instalar do Model Explorer é executado localmente no dispositivo, no Colab ou em um arquivo do Python.

Execução local no dispositivo

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

Esses comandos inicializarão um servidor em localhost:8080 e abrirão o app da Web do Model Explorer em uma guia de navegador. Saiba mais sobre a utilização de linhas de comando do Model Explorer no guia de linhas de comando.

Quando você tiver o servidor localhost em execução, faça upload do arquivo de modelo de seu computador (os formatos com suporte incluem aqueles usados pelo JAX, PyTorch, TensorFlow e TensorFlow Lite) e, se necessário, selecione o melhor adaptador para o modelo em questão usando o menu suspenso "Adapter". Clique aqui para saber como utilizar o sistema de extensão de adaptador do Model Explorer para visualizar formatos de modelos sem suporte.

Execução em um notebook do 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)

Após a execução da célula, o Model Explorer é exibido em um iframe incorporado em uma nova célula. No Chrome, a IU também mostra um botão "Open in new tab", no qual você pode clicar para mostrar a IU em uma guia separada. Clique aqui para saber mais sobre a execução do Model Explorer no Colab.

Visualize models via the Model Explorer API

Abra o visualizador usando as APIs do Model Explorer

O pacote model_explorer fornece APIs convenientes para permitir a visualização de modelos de arquivos ou de um módulo do PyTorch, além de uma API de nível mais baixo para visualizar modelos de várias origens. Não deixe de instalá-lo primeiro, acompanhando o guia de instalação. Para saber mais, confira o guia de APIs do Model Explorer.

Veja abaixo um exemplo de como visualizar um modelo do PyTorch. A visualização de modelos do PyTorch requer uma abordagem ligeiramente diferente de outros formatos, porque o PyTorch não tem um formato de serialização padrão. O Model Explorer oferece uma API especializada para visualizar modelos do PyTorch diretamente, usando o ExportedProgram de 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)

Independentemente de como os modelos sejam visualizados, nos bastidores, o Model Explorer implementa a renderização de grafos acelerada por GPU com o WebGL e o three.js, que habilita uma experiência de visualização fluida de 60 QPS, mesmo em grafos que contenham dezenas de milhares de nós. Caso tenha interesse em saber mais sobre como o Model Explorer renderiza grafos grandes, consulte o blog Google Research,

Noções básicas sobre a arquitetura do modelo

Os modelos grandes podem ser complexos, mas o Model Explorer os torna mais fáceis de entender dividindo a visualização em camadas hierárquicas. Dê uma olhada no modelo MobileBert mostrado abaixo: fica claro como a máscara de autoatenção e a incorporação são alimentadas em uma camada transformer. Você pode até se aprofundar na camada embedding para entender as relações entre os diferentes tipos de incorporação. A visualização hierárquica do Model Explorer torna mais fáceis de entender até mesmo as arquiteturas de modelo mais complexas.

Com informações de camada, como "self_attention_mask", "embedding" e "transformer_layer", fica fácil entender a arquitetura do modelo.

Depuração de erros de conversão

Converter modelos de um formato para outro (como do PyTorch para TFLite) pode ser complicado, mas o Model Explorer ajuda a comparar os grafos originais e convertidos lado a lado. Isso simplifica a identificação de quaisquer alterações que possam afetar o desempenho do modelo. Por exemplo, nas imagens abaixo, você pode ver como um subgrafo dentro de uma camada mudou durante a conversão, o que ajuda a identificar e corrigir possíveis erros.

A side-by-side comparison of the original PyTorch model (left) and the optimized TensorFlow Lite model (right)

Uma comparação lado a lado do modelo do PyTorch original (esquerda) e do modelo do TensorFlow Lite otimizado (direita), que é útil para depurar a conversão de grafos de modelo.

Desempenho da depuração e acurácia numérica com sobreposições de dados por nó

Um recurso avançado do Model Explorer é a capacidade de sobrepor dados por nó em um grafo, permitindo classificar, pesquisar e estilizar nós usando os valores nesses dados. Combinado com o sistema de visualização hierárquica do Model Explorer, esse recurso permite identificar rapidamente o desempenho ou os gargalos numéricos. O exemplo abaixo mostra o erro quadrático médio em cada nó entre um modelo quantizado do TFLite e sua contraparte de ponto flutuante. O Model Explorer destaca que a queda de qualidade está próxima da parte inferior do grafo, fornecendo as informações necessárias para ajustar o método de quantização. Para saber mais sobre como trabalhar com dados personalizados no Model Explorer, confira nossa documentação detalhada no Github.

Essa sobreposição de dados por nó permite que os usuários identifiquem rapidamente problemas numéricos ou de desempenho em um modelo.

Utilização no Google

Entre os usuários mais proeminentes do Model Explorer no Google estão o Waymo e o Google Silicon. O Model Explorer desempenhou um papel crucial para ajudar essas equipes a depurar e otimizar modelos no dispositivo, como o 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

O que vem a seguir?

Para nós, este é apenas o começo. Nos próximos meses, nos concentraremos em aprimorar o núcleo, refinando os principais recursos da IU, como diferenciação e edição de grafos, e em melhorar a capacidade de extensão para permitir que você integre suas próprias ferramentas ao Model Explorer.

Agradecimentos

_{Este trabalho é uma colaboração entre várias equipes funcionais do Google. Gostaríamos de agradecer aos engenheiros 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 e Ruofei Du, à nossa Gerente técnica de programa, Kristen Wright, e a nosso Gerente de produto, Aaron Karp. Nossos agradecimentos também para a equipe de UX, inclusive Zi Yuan, Anila Alexander, Elaine Thai, Joe Moran e Amber Heinbockel.}