モデル エクスプローラについてお知らせします。モデル エクスプローラは、強力なグラフ視覚化ツールです。モデルのアーキテクチャや、ノード間のデータフローを視覚化する機能は、独自のモデルの開発、フォーマット間のモデルの変換、特定のデバイスへの最適化、パフォーマンスや品質のデバッグなどに非常に便利です。大規模なグラフでも直感的で階層的な視覚化を行えるので、大規模モデルの複雑な作業に対応でき、特にモデルをエッジデバイス用に最適化する際に役立ちます。
本記事は、Google AI Edge のデベロッパー リリースを紹介するシリーズの 3 回目となるブログ投稿です。最初の 2 回では、オンデバイスで PyTorch モデルと高パフォーマンス LLM を実現する AI Edge Torch と Generative API を紹介しました。
モデル エクスプローラは、もともと Google の研究者やエンジニア向けのユーティリティとして開発されましたが、現在は Google AI Edge プロダクト ファミリーの 1 つとして一般公開されています。モデル エクスプローラの初期バージョンでは、以下の機能が提供されます。
このブログ投稿では、モデル エクスプローラの使い方のほか、いくつかの強力なユースケースについて説明します。詳しいドキュメントやサンプルは、こちらをご覧ください。
モデル エクスプローラは、簡単にインストールできる 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] ドロップダウン メニューからモデルに最適なアダプタを選びます。モデル エクスプローラのアダプタ拡張システムを使ってサポートされていない形式のモデルを視覚化する方法は、こちらをクリックしてください。
# モデルをダウンロード（この例では 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 で実行する際の詳細については、こちらをご覧ください。
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 モデルをご覧ください。セルフアテンション マスクと埋め込みがトランスフォーマー レイヤにどのようにフィードされるかがよくわかります。埋め込みレイヤに注目して、さまざまなタイプの埋め込みの関係を理解することもできます。どんなに複雑なモデル アーキテクチャであっても、モデル エクスプローラの階層ビューならわかりやすく表示できます。
モデルをあるフォーマットから別のフォーマット（PyTorch から TFLite など）に変換するのは、簡単なことではないかもしれません。しかし、モデル エクスプローラなら、元のグラフと変換されたグラフを並べて比較できるので、モデルのパフォーマンスに影響を与えそうな変更点を簡単に見つけることができます。たとえば、次の画像では、変換によってレイヤ内のサブグラフがどのように変化したかを確認できます。これは潜在的なエラーの特定や修正に役立ちます。
モデル エクスプローラには、ノードごとのデータをグラフにオーバーレイする強力な機能があります。このデータ値を使ってノードを並べ替えたり、検索したり、スタイルを設定したりすることができます。これを階層ビューシステムと組み合わせると、ノードごとのデータ オーバーレイから、パフォーマンスや数値のボトルネックをすばやく見つけることができます。次の例は、量子化した TFLite モデルと、浮動小数点数版のモデルの各ノードでの平均二乗誤差を示しています。モデル エクスプローラを使うと、品質の低下がグラフの下部近くで起きていることがわかるので、量子化の方法の調整に必要な情報を得られます。モデル エクスプローラでのカスタムデータの操作については、Github の詳細ドキュメントをご覧ください。
Google で特によく知られているモデル エクスプローラのユーザーは、Waymo と Google Silicon です。モデル エクスプローラは、こういったチームが Gemini Nano のようなオンデバイス モデルのデバッグと最適化を行ううえで重要な役割を果たしました。
これはまだほんの始まりに過ぎません。今後数か月で、グラフの差分や編集などの主要な 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 チームにも感謝を捧げます。
