モデルエクスプローラ: エッジデバイスの ML モデルをシンプルに

6月 26, 2024

Kristen Wright Technical Program Manager

Eric Yang Software Engineer

モデルエクスプローラについてお知らせします。モデルエクスプローラは、強力なグラフ視覚化ツールです。モデルのアーキテクチャや、ノード間のデータフローを視覚化する機能は、独自のモデルの開発、フォーマット間のモデルの変換、特定のデバイスへの最適化、パフォーマンスや品質のデバッグなどに非常に便利です。大規模なグラフでも直感的で階層的な視覚化を行えるので、大規模モデルの複雑な作業に対応でき、特にモデルをエッジデバイス用に最適化する際に役立ちます。

本記事は、Google AI Edge のデベロッパーリリースを紹介するシリーズの 3 回目となるブログ投稿です。最初の 2 回では、オンデバイスで PyTorch モデルと高パフォーマンス LLM を実現する AI Edge Torch と Generative API を紹介しました。

モデルエクスプローラは、もともと Google の研究者やエンジニア向けのユーティリティとして開発されましたが、現在は Google AI Edge プロダクトファミリーの 1 つとして一般公開されています。モデルエクスプローラの初期バージョンでは、以下の機能が提供されます。

数万規模のノードを持つ非常に大規模なモデルのグラフでも、60 FPS のスムーズな視覚化を実現

JAX、PyTorch、TensorFlow、TensorFlow Lite などの一般的な ML フレームワークをサポート

Colab ノートブックで直接実行

多くのモデルフォーマットへの拡張性を提供するカスタム拡張機能

メタデータ（属性、入出力など）やカスタムデータ（パフォーマンスなど）をノードにオーバーレイで直接表示

グラフを並べて比較できるスプリットペインビューなど、作業を高速化する強力な UI 機能

このブログ投稿では、モデルエクスプローラの使い方のほか、いくつかの強力なユースケースについて説明します。詳しいドキュメントやサンプルは、こちらをご覧ください。

スタートガイド

モデルエクスプローラは、簡単にインストールできる PyPI パッケージです。デバイスのローカルでも、Colab でも、Python ファイルでも実行できます。

デバイスでローカル実行する

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

このコマンドは、localhost:8080 でサーバーを起動し、ブラウザのタブでモデルエクスプローラのウェブアプリを開きます。モデルエクスプローラのコマンドラインの使い方については、コマンドラインガイドをご覧ください。

ローカルホストでサーバーを実行したら、コンピュータからモデルのファイルをアップロードし（サポートされている形式は、JAX、PyTorch、TensorFlow、TensorFlow Lite などです）、必要に応じて [Adapter] ドロップダウンメニューからモデルに最適なアダプタを選びます。モデルエクスプローラのアダプタ拡張システムを使ってサポートされていない形式のモデルを視覚化する方法は、こちらをクリックしてください。

Colab ノートブックで実行する

# モデルをダウンロード（この例では 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)
 
# モデル エクスプローラをインストール
pip install ai-edge-model-explorer
 
# ダウンロードした EfficientDet モデルを視覚化
import model_explorer
model_explorer.visualize(model_path)

セルを実行すると、新しいセルに埋め込まれた iFrame にモデルエクスプローラが表示されます。Chrome では、新しいタブで開くためのボタンも UI に表示されるので、これをクリックして別のタブに UI を表示できます。モデルエクスプローラを Colab で実行する際の詳細については、こちらをご覧ください。

Visualize models via the Model Explorer API

モデルエクスプローラ API を使って視覚化ツールを開く

model_explorer パッケージでは、ファイルや PyTorch モジュールのモデルを視覚化する便利な API と、複数のソースからモデルを視覚化する低レベル API が提供されます。最初に、インストールガイドに従ってこのパッケージをインストールしてください。詳細については、モデルエクスプローラ API ガイドをご覧ください。

以下は、PyTorch モデルを視覚化する方法の例です。PyTorch には標準的なシリアル化形式が存在しないため、PyTorch モデルは別の形式とわずかに異なるアプローチで視覚化する必要があります。モデルエクスプローラには、torch.export.export の ExportedProgram を使って PyTorch モデルを直接視覚化する特別な API があります。

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)

モデルエクスプローラは、モデルを視覚化する方法に関係なく、内部的に WebGL と three.js の GPU アクセラレーションを利用してグラフをレンダリングします。そのため、数万規模のノードを持つグラフでも、60 FPS のスムーズな視覚化を実現できます。モデルエクスプローラが大規模なグラフをレンダリングする方法について詳しく知りたい方は、Google Research ブログをお読みください。

モデルのアーキテクチャを理解する

大規模モデルには複雑なものもありますが、モデルエクスプローラは階層的なレイヤに分割して視覚化することで、理解しやすくしています。下の図の MobileBert モデルをご覧ください。セルフアテンションマスクと埋め込みがトランスフォーマーレイヤにどのようにフィードされるかがよくわかります。埋め込みレイヤに注目して、さまざまなタイプの埋め込みの関係を理解することもできます。どんなに複雑なモデルアーキテクチャであっても、モデルエクスプローラの階層ビューならわかりやすく表示できます。

「self_attention_mask」、「embedding」、「transformer_layer」などのレイヤ情報が表示されるので、モデルのアーキテクチャを簡単に理解できる。

変換エラーのデバッグ

モデルをあるフォーマットから別のフォーマット（PyTorch から TFLite など）に変換するのは、簡単なことではないかもしれません。しかし、モデルエクスプローラなら、元のグラフと変換されたグラフを並べて比較できるので、モデルのパフォーマンスに影響を与えそうな変更点を簡単に見つけることができます。たとえば、次の画像では、変換によってレイヤ内のサブグラフがどのように変化したかを確認できます。これは潜在的なエラーの特定や修正に役立ちます。

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

オリジナルの PyTorch モデル（左）と最適化された TensorFlow Lite モデル（右）を並べて比較 – モデルグラフ変換のデバッグに有効。

ノードごとのデータオーバーレイでパフォーマンスと数値精度をデバッグする

モデルエクスプローラには、ノードごとのデータをグラフにオーバーレイする強力な機能があります。このデータ値を使ってノードを並べ替えたり、検索したり、スタイルを設定したりすることができます。これを階層ビューシステムと組み合わせると、ノードごとのデータオーバーレイから、パフォーマンスや数値のボトルネックをすばやく見つけることができます。次の例は、量子化した TFLite モデルと、浮動小数点数版のモデルの各ノードでの平均二乗誤差を示しています。モデルエクスプローラを使うと、品質の低下がグラフの下部近くで起きていることがわかるので、量子化の方法の調整に必要な情報を得られます。モデルエクスプローラでのカスタムデータの操作については、Github の詳細ドキュメントをご覧ください。

このノードごとのデータオーバーレイから、モデルで発生しているパフォーマンスや数値の問題をすばやく特定可能。

Google での使用事例

Google で特によく知られているモデルエクスプローラのユーザーは、Waymo と Google Silicon です。モデルエクスプローラは、こういったチームが 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

次のステップ

これはまだほんの始まりに過ぎません。今後数か月で、グラフの差分や編集などの主要な UI 機能を改良すること、および独自のツールをモデルエクスプローラに組み込めるようにして拡張性を向上させることにより、中核機能を重点的に強化する予定です。

謝辞

_{これは、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 をはじめとする UX チームにも感謝を捧げます。}