Google AI Edge Torch: モバイルデバイス向けPyTorchモデルの高性能推論を実現する
2024年、Googleは「AI Edge Torch」のベータ版を発表しました。これは、PyTorchからTensorFlow Lite (TFLite) ランタイムへの直接的な道を提供し、モデルのカバレッジとCPU性能を向上させることを目的としています。これにより、PyTorchを利用する開発者の皆さんは、よりスムーズにモバイルデバイス向けのMLパイプラインを構築できるようになります。
AI Edge Torchの特徴
- **PyTorchとの直接統合**: PyTorchユーザーにとって、使いやすいAPIを提供。
- **優れたCPU性能と初期のGPUサポート**: モバイルデバイス向けの用途に最適化されています。
- **70以上のモデルに対応**: torchvision、timm、torchaudio、HuggingFaceのモデルを検証済み。
- **TFLiteランタイムとの互換性**: デプロイコードの変更が不要です。
- **モデルエクスプローラーの視覚化**: ワークフローの複数の段階におけるモデルの視覚化が可能。
使い方のシンプルさ
AI Edge Torchは、PyTorchコミュニティに最適な体験を提供するために設計されています。以下に、実際の使用例を示します。
import torchvision
import ai_edge_torch
# モデルの初期化
resnet18 = torchvision.models.resnet18().eval()
# サンプル入力の定義
sample_input = (torch.randn(4, 3, 224, 224),)
# モデルの変換
edge_model = ai_edge_torch.convert(resnet18, sample_input)
# Python内での推論実行
output = edge_model(*sample_input)
# デバイス向けにTfLiteモデルをエクスポート
edge_model.export('resnet.tflite')
このコードにより、PyTorchモデルが簡単にTFLiteフォーマットに変換され、モバイルデバイス上で実行できるようになります。
カバレッジとパフォーマンス
これまで多くの開発者が、PyTorchモデルをTFLiteで動かすためにONNX2TFなどのコミュニティ提供のパスを利用していましたが、AI Edge Torchの開発目標は、開発者の負担を軽減し、優れたモデルカバレッジを提供することです。
- カバレッジ: 定義されたモデルセットに対して、既存のワークフローに比べて大幅な改善を示しました。
- パフォーマンス: ONNX2TFのベースラインと比較して、一貫した性能を実現。特に、Pixel 8での推論レイテンシが向上しています。
初期の導入とパートナーシップ
AI Edge Torchの開発においては、Shopify、Adobe、Nianticといった早期採用パートナーとの密接な協力がありました。Shopifyでは、製品画像の背景除去を行うために、AI Edge Torchが実装されています。
今後の展開
今後の数ヶ月で、モデルカバレッジの拡大、GPUサポートの改善、さらなる量子化モードの追加を行い、1.0リリースに向けた進展を続ける予定です。また、次回の記事では、カスタムGenAIモデルをエッジで動作させるためのAI Edge Torch Generative APIについて詳しく解説します。
感謝の意
最後に、初期アクセス顧客やハードウェアパートナー、XNNPACKのエコシステム貢献者に感謝の意を表します。彼らのフィードバックにより、初期のバグを発見し、スムーズな開発者エクスペリエンスを確保できました。
AIモデルのデプロイは複雑ですが、AI Edge Torchにより、開発者は容易にアプローチできるようになります。今後の展開に期待しましょう。
