Zenn
🧜‍♂️

[Triton-Inference-Server] なぜTorchScriptはmodel.ptを置くだけで良いのか

2024/10/20に公開
Python
PyTorch
MLOps
triton
torchscript
tech

はじめに

NVIDIAのTriton Inference Serverは、さまざまな機械学習モデルを効率的にデプロイし、スケーラブルな推論サービスを提供するための強力なツールです。特に、PyTorchで開発されたモデルをTorchScript形式に変換することで、Tritonにおいてmodel.ptファイルを所定の場所に配置するだけでモデルをデプロイできます。

本記事では、なぜTorchScriptモデルがmodel.ptを置くだけで動作するのか、その背後にある理由とTritonの仕組みについて詳しく解説します。

TL;DR

  • この記事は、NVIDIAのTriton Inference ServerがPyTorchのTorchScriptモデルをmodel.ptファイルを所定のディレクトリに置くだけで動作させる理由を説明しています。
  • TorchScriptモデルは自己完結型であり、モデルの計算グラフと学習済みパラメータが一つのファイルにまとめられているため、Pythonの依存なしにロード・実行できます。
  • Triton Inference Serverは、モデルのファイル形式に応じて自動的に適切な方法でモデルをロードする柔軟な仕組みを持っています。
  • 特定のディレクトリ構造（モデルリポジトリ）に従ってmodel.ptを配置することで、Tritonはモデルを自動検出し、追加の設定やコードなしでデプロイできます。
  • これにより、高性能でスケーラブルな推論サービスの提供が簡単になります。

1. TorchScriptモデルとは

1.1 TorchScriptの概要

TorchScriptは、PyTorchの一部で、PyTorchで作成したモデルを保存して後で再利用できるようにする機能を提供しています。TorchScriptを使用すると、PyTorchのモデルをPythonに依存しない形式で表現でき、以下の特徴があります。

  • 自己完結型：モデルのアーキテクチャ（計算グラフ）と学習済みのパラメータが一つのファイルにまとめられる。
  • 言語非依存性：Python環境がなくても、C++など他の言語やランタイムでモデルをロードして実行可能。
  • JITコンパイル：Just-In-Timeコンパイルにより、推論時のパフォーマンスが向上する。

1.2 TorchScriptモデルの作成

PyTorchで学習したモデルをTorchScript形式に変換するのは非常に簡単です。以下のように、torch.jit.traceやtorch.jit.scriptを使用してモデルを保存できます。

pythonimport torch

# モデルの定義と学習済みパラメータのロード
model = ...  # あなたのモデル
model.load_state_dict(torch.load('model_weights.pth'))
model.eval()

# TorchScriptモデルへの変換
scripted_model = torch.jit.trace(model, example_input)
scripted_model.save('model.pt')

Pytorchのモデル保存方法については、以下の記事を参考にしてもらうと良いと思います。

https://takaherox.hatenablog.com/entry/2021/01/09/230332

2. Triton Inference Serverのアーキテクチャ

2.1 モジュラー設計

Triton Inference Serverは、プラグイン可能なバックエンドアーキテクチャを持ち、さまざまなフレームワークやモデル形式をサポートしています。主な特徴は以下の通りです。

  • マルチフレームワーク対応：TensorFlow、PyTorch、ONNX、TensorRTなど、多数のフレームワークをサポート。
  • バックエンドプラグイン：各フレームワークに対応するバックエンドを持ち、モデルのロードや推論を担当。
  • モデルリポジトリの自動検出：特定のディレクトリ構造に従ってモデルを配置することで、Tritonが自動的にモデルを検出し管理。

2.2 PyTorchバックエンド

TritonはPyTorch用の専用バックエンドを持っており、TorchScript形式のモデルを直接ロードして推論を実行できます。このバックエンドは以下の機能を提供します。

  • モデルのロード：model.ptファイルからTorchScriptモデルをロード。
  • 入力・出力の管理：モデルの入出力テンソルを適切に処理。
  • 推論の最適化：バッチ処理や並列処理によるパフォーマンス向上。

3. TorchScriptモデルが動作する仕組み

3.1 自己完結型のモデル

TorchScriptモデルは、モデルの計算グラフと学習済みのパラメータを一つのmodel.ptファイルにまとめています。 これにより、モデルのロードには追加のPythonコードや依存関係が不要となります。

3.2 Tritonのモデルロードプロセス

  1. モデルリポジトリのスキャン：Tritonは起動時にモデルリポジトリをスキャンし、利用可能なモデルを検出します。
  2. バックエンドの選択：モデルの拡張子（.pt）から適切なバックエンド（この場合はPyTorchバックエンド）を選択します。
  3. モデルのロード：PyTorchバックエンドがmodel.ptファイルをロードし、メモリ上にモデルを展開します。
  4. 入出力の設定：config.pbtxtファイル（必要に応じて）からモデルの入出力情報を取得します。

3.3 追加の設定が不要

TorchScriptモデルは自己完結しているため、model.ptファイルを所定のディレクトリに配置するだけで、Tritonはモデルを適切にロードし、推論を実行できます。追加のPythonスクリプトや設定は基本的に不要です。

4. モデルの配置と自動検出

4.1 モデルリポジトリの構造

Tritonがモデルを自動的に検出するためには、特定のディレクトリ構造に従う必要があります。TorchScriptモデルの場合、以下のようになります。　参照

models/
└── your_model_name/
    ├── 1/
    │   └── model.pt        # TorchScriptモデルファイル
    └── config.pbtxt        # モデルの設定ファイル（必要に応じて）
  • your_model_name：モデルの名前で、API経由でのリクエスト時に使用します。
  • 1：モデルのバージョン番号。Tritonはバージョン管理をサポートしています。
  • model.pt：TorchScript形式のモデルファイル。

4.2 Tritonによる自動検出とデプロイ

  • Tritonはモデルリポジトリを監視し、新しいモデルや更新されたモデルを自動的にロードします。
  • モデル名、バージョン、ファイル構造に基づいて、適切にモデルを管理します。

5. まとめ

Triton Inference Serverでは、TorchScriptモデルの自己完結性とPyTorchバックエンドのサポートにより、model.ptファイルを適切なディレクトリ構造に配置するだけで、モデルをデプロイできます。これは、TorchScriptがモデルの定義とパラメータを一つのファイルにまとめているため、追加のコードや依存関係が不要であることが理由です。

また、Tritonのアーキテクチャにより、モデルリポジトリからの自動検出・ロードが可能となり、煩雑なデプロイ作業を大幅に簡素化できます。これにより、高性能でスケーラブルな推論サービスの提供が容易になります。

参考資料

この記事がお役に立てば幸いです。ご質問やフィードバックがございましたら、お気軽にご連絡ください。

Discussion