コンテンツにスキップ

Triton Model Navigator で PyTorch, TensorFlow, ONNX モデルを最適化

概要

Triton Model Navigatorは、NVIDIA GPU を活用したディープラーニングモデルの最適化とデプロイを目的とする推論ツールキットだ。PyTorch、TensorFlow、ONNX で実装したモデルとパイプラインを対象に、TensorRT への移行を効率化する。

主な機能

Triton Model Navigator は、次の機能を備えている。

  • 最適化の自動化: モデルのエクスポート、変換、検証、プロファイリングなどの手順を自動化。
  • フレームワーク間互換性: PyTorch、TensorFlow、ONNX など、さまざまな機械学習フレームワークをサポートし、各フレームワークの最適化機能を活用して、最適なデプロイオプションを効率的に探索。
  • パフォーマンス向上: ResNet や BERT などのモデルのパフォーマンスを向上させ、推論を効率化。
  • パイプラインの合理化: Stable Diffusion や Whisper などのモデルの Python コードパイプラインを、PyTorch 専用の Inplace Optimization を活用して効率化。
  • 変換の自動化: TensorRT と Torch-TensorRT に焦点を当て、さまざまなフォーマット間でのモデルのエクスポートと変換プロセスを自動化。
  • 精度の検証: 変換後モデルが元モデルに対し正しい出力を生成することを確認。
  • プロファイリング: 応答までにかかる時間（レイテンシ）や単位時間あたりの処理量（スループット）などの指標に基づいて最適なフォーマットを選択し、ハードウェアの使用率を最適化。
  • デプロイの自動化: 専用 API を通じて、PyTriton と Triton Inference Server へのモデルとパイプラインのデプロイを自動化。
  • QAT サポート: 量子化による精度変化を学習中に織り込む QAT（Quantization Aware Training）に対応し、INT8 の形式でモデルを最適化。

setup

torchの extra dependencies を含めてインストールする場合は、次のコマンドを実行する。ほかにもtensorflowやjaxなどの extra dependencies があるので、使うフレームワークに合わせて指定する。

1
uv pip install -U --extra-index-url https://pypi.ngc.nvidia.com "triton-model-navigator[torch]"

最適化の流れ

v0.12.0 のQuick Startでは、モデルと代表的な入力データを渡して最適化し、変換した候補を検証・計測してから実行する流れになっている。PyTorch、TensorFlow、ONNX それぞれの API で、エクスポートや最適化を行う。最適化後のモデルは、PyTriton や Triton Inference Server へ配置して利用できる。

この記事ではモデルごとの実装手順ではなく、導入時に試した環境と回避策を扱う。フレームワーク別のコード例は Quick Start を参照してほしい。

tips

  • torch のモデルを実際に最適化すると、model navigator 内部での ONNX への export がときどき失敗した。その場合は、torch のモデルをあらかじめ ONNX へ export してから model navigator を実行すると、最適化が円滑に進む。
  • v0.12.0 時点では Python 3.11 で実行するとエラーが発生したため、Python 3.10 で実行することで解消できる。
  • v0.12.0 時点では、polygraphy のバージョンと相性が悪く、アップデートが必要だった。次のコマンドでバージョンを更新すると解決できる。
1
uv pip install --index-strategy unsafe-best-match polygraphy==0.49.* --extra-index-url https://pypi.ngc.nvidia.com

類似のツール

speedsterやBigDL-Nanoも、同様の最適化を目的としている。ただし、いずれも開発が止まっているように見える。 実際に両ツールの最新バージョンを使ってみると、エラーや依存関係の問題が発生する。 2024/10/28 時点で Triton Model Navigator は更新が続いているため、Triton Model Navigator の利用をお勧めする。

まとめ

Triton Model Navigator は、PyTorch、TensorFlow、ONNX モデルの最適化と TensorRT への変換を自動化するツールキットだ。精度検証やプロファイリングも含めて、デプロイまで一貫して扱える。ただし、筆者が試した範囲では、次の回避策が必要になる場面があった。

  • ONNX への事前 export
  • Python バージョンの指定
  • 依存パッケージのバージョン調整

類似ツールの開発が停滞しているように見える一方で、Triton Model Navigator は 2024/10/28 時点でも更新が続いている。長期的に採用するうえでは、継続的な更新が安心材料になる。