Triton Inference Server のバッチ サイズを紐解く

みなさま Medium ではご無沙汰しております。NVIDIA の山崎です。同時に公開した Triton Inference Server 2022 年 6 月のリリース概要はご覧いただけましたでしょうか。この記事の中で、このバージョンから自動生成される設定ファイルのデフォルト挙動が変わる旨について言及しました。変更内容としては、dyamic batching がデフォルトで有効になるというものですが、一方でデプロイ可能なモデルに何か制約がつくのか? そもそも Triton のバッチ サイズ関連の設定ってどうなってるの? など、疑問を持たれているかもしれません。この記事では、そうした疑問を少しでも解消すべく、Triton におけるバッチ サイズの設定について詳細な説明を試みます。

長くて読む気がしない人に贈るバッチ サイズ周りの設定ポイント

• 設定のことなんて考えたくない
  → 冒頭で言及している設定の自動生成機能を利用 (=Triton 起動時に --strict-model-config=false オプションを指定)
• 手っ取り早く任意のバッチ サイズを扱いたい
  → モデルのエクスポート時に可変バッチ (=バッチ次元に -1) を設定し、 max_batch_size は指定せず、dims の先頭を -1 にする
• 複数のリクエストからバッチを構成したい
  → やはりモデルのエクスポート時、バッチ サイズが可変になるよう設定しておき、 max_batch_size に想定している最大バッチ サイズを設定し、dims からバッチ次元を除外してテンソル形状を定義し、dynamic_batchingを設定する

なお上記、後述しますがバックエンドによって (特に TorchScript 利用時は) 条件が緩和される方向の例外があったりするため、最終的には多少実験していただくことを推奨します。

Triton で指定可能なバッチ サイズ関連の設定

そもそも「バッチ サイズ」といえば、多くの場合、複数のデータをまとめて処理する際のデータの個数を指し、入力されるテンソルの最初の次元で扱われるものです。画像認識であればまとめて処理される画像の枚数、文書分類であれば分類対象のドキュメントの数、といった具合です。こうしたまとめて実行されるデータの数は、オフライン推論であれば固定することが比較的容易である一方、オンライン推論では様々なサイズを考慮する必要があり、バッチ サイズを固定することは好ましくない場合があります。

こうした状況をサポートするべく、Triton には様々な設定が存在しています。記事執筆時点で は、バッチ関連の主な設定として以下のものがあります。

• max_batch_size
• dims
• dynamic_batching の preferred_batch_size

max_batch_size には、許容する最大のバッチ サイズを 1 以上の値で指定します。ゼロ、もしくは指定しないことも可能です。

dims には、入出力テンソルの形状を指定します。-1 を指定すると、その次元は可変長となります。バッチ サイズという点では、主に先頭要素を -1 にするかどうかがポイントとなります。また詳細は後述しますが、max_batch_size を指定する場合、dims の先頭要素はバッチの次元ではなくなるケースがほとんどです。

preferred_batch_size は、後述する dynamic batching を利用する際のバッチ サイズを制御するための設定です。主に TensorRT で複数のバッチ サイズ向けに最適化されている状況で利用されることが想定されており、特定のバッチ サイズが特に性能向上に寄与する場合に指定することが推奨されます。これが設定されている場合、Triton のスケジューラは指定されている中から最大のサイズでバッチを構成し、推論実行しようとします。なお、この設定が指定されていない場合は、max_batch_size 以下の最大サイズでバッチを構成して推論します。

またこれらに加えて、実際にはモデル保存時の設定も関連します。たとえば、TensorFlow 利用時の SavedModel であれば tf.TensorSpec などを利用した signature の指定であったり、PyTorch を ONNX にエクスポートする場合であれば torch.onnx.export に渡す args や dynamic_axes などです。

このように、バッチ サイズに関する設定には主に max_batch_size、dims、dynamic_batching の preferred_batch_size、およびモデル保存時の設定の 4 つの要素が存在しています。ではこれらはどのように関連しているのでしょうか。

max_batch_size と dims: default scheduler の場合

まず max_batch_size と dims についてです。何をおいても押さえておくべきポイントとして、「max_batch_size 未指定もしくは 0 が指定される場合、デフォルトでは、クライアントから送信できるテンソルは dims と一致する形状のみ」というものがあります。たとえば画像分類モデルなどで、(224, 224, 3) というテンソルを入力として受け取るモデルを考えます。このとき、Triton 側に以下のような設定を入れた場合、

max_batch_size: 0
...
input [{
...
    dims: [224, 224, 3]
...

複数の画像をまとめて送信することはできず、Triton 側での推論は常にバッチ サイズ 1 相当の状況で実行されることになります。加えて、モデルのエクスポート時に一次元目をバッチの次元として指定していた場合、上記の設定ではモデルロード自体に失敗します。(注: 後述しますがバックエンドによっては通るものもあります)

では dims にバッチの次元に相当する値を入れる場合どうなるでしょうか。この場合も原則として、クライアント側のリクエストには dims と一致する形状のテンソルのみが許容されます。一方で先に述べたとおり、-1 を指定し可変長として扱うことで、任意のバッチ サイズのリクエストを送信することはできます。たとえば以下のような設定です。

...
input [{
...
    dims: [-1, 224, 224, 3]
...

これにより、クライアント側で複数データを扱う必要のある状況はカバーされます。ただし、この指定を行うためには、モデルのエクスポート時に先頭次元が -1 で可変長を仮定している必要があります。(注: こちらもバックエンドによっては、エクスポート時に固定サイズを前提としていても問題なく動くケースがあります)

とはいえ実利用を考えた場合、上記の設定では問題になるケースがあると想像できます。特に指定をしない場合、Triton では default scheduler が利用されます。ここまで述べてきたような max_batch_size を設定しないケースでは、この default scheduler が利用され、リクエストごとに推論実行されます。しかし Triton を利用するような状況では、多数のクライアントが独立してリクエストを送信してくることがほとんどで、リクエストごとに推論実行するのはあまり効率的ではありません。そこで、Triton には性能向上のための重要な機構として、dynamic batcher が存在します。

max_batch_size と dims: dynamic batcher の場合

詳細はこの記事の範囲から外れるため割愛しますが、dynamic batcher は複数のリクエストから動的にバッチを構成し、複数のリクエストを束ねて推論実行する機構です。有効化するためには、設定に dynamic_batching: {}という記述を追加する必要があります。このとき注意すべきポイントとして、max_batch_size が未指定のままだとリクエストごとに推論する動作は変わらない、というものがあります。

そこで通常は max_batch_size に 1 以上の値を指定し、dynamic batching が機能するように設定することが推奨されます。注意すべきは、dims の先頭にはバッチ次元を指定「しない」という点です。max_batch_size を指定すると、Triton 側では、最終的な推論時の入力テンソルの形状を [-1] + dims として扱います。そのため dims の先頭にバッチ サイズ相当の値が入っている場合、Triton は、余分な次元が一つ増えたテンソルが入力されると仮定して処理を進め、エラーとなります。たとえば先の例であれば、以下のように設定を変更することになります。

max_batch_size: 32
...
input [{
...
    dims: [224, 224, 3]
...

ほとんど最初の例と同じですが、max_batch_size の値だけが異なります。これにより、クライアントからは任意のサイズのバッチを送信できるようになります。なおこの場合も、モデルをエクスポートする際、可変バッチ サイズに対応するように指定しておく必要があります。(注: これもバックエンドによります)

さらに加えて dynamic_batching を指定することで、複数のリクエストをひとつに束ねて推論実行できるようになります。例えば以下のように設定を追加します。

dynamic_batching: {}

ただしこの場合デフォルトでは待ち時間が 0 のため、運良く近いタイミングで複数リクエストが同時にサーバーへ到着するか、推論の実行時間が長いことで多数のリクエストがキューに滞留しない限り、シーケンシャルに推論が実行されます。確実に複数のリクエストを一つのバッチに束ねるには、max_queue_delay_microseconds をほどよい値に設定することが重要となります。例えば以下のような値です。

dynamic_batching: {
  # Wait for up to 10 msec.
  max_queue_delay_microseconds: 10000
}

実際には Performance Analyzer や Model Analyzer などを利用して調整することが推奨されます。こうしたチューニングの指針は Recommended Configuration Process などを参照ください。また、実際に構成されたバッチ サイズを知りたい場合、起動時オプションの --log-verbose=1 を指定して出力されるログから拾う方法と、metrics を利用する方法があります。詳細は割愛しますが、”Inference Count (nv_inference_count)” と “Execution Count (nv_inference_exec_count)” を組み合わせることで、実際に実行された平均バッチ サイズを計算できます。

preferred_batch_size を考慮する必要のある場合

最後に preferred_batch_size です。前述した通り、TensorRT のように特定のバッチ サイズで特に強く最適化されてる、という場合に考慮する必要のあるパラメータです。基本的には max_batch_size に類似の設定とみなしてよく、例えば以下のように記述すると、バッチ サイズ 1 もしくは 2 で推論実行しようと試みます。

dynamic_batching: {
  preferred_batch_size: [1, 2]
}

なお、特に TensorRT を利用する場合、固定バッチ サイズを前提とするモデルでも問題なく動作する「ことがある」点には注意が必要です。preferred_batch_size を利用する場合に特に発生しやすくなると予想されますが、たとえばバッチ サイズ 2 を前提とするモデルに対して、preferred_batch_size: [1] を設定する状況を考えます。このとき、複数のクライアントからバッチ サイズ 1 でリクエストが送信されると、たまたま同時に 2 つ以上のリクエストが到着した場合には、バッチ サイズ 2 で推論実行されるため成功しますが、多くの場合バッチ サイズ 1 で推論実行が走ります。しかしながら、モデルが前提とするバッチ サイズと異なるため、サイズの不一致でエラーが発生してしまい、多数のリクエストが失敗することになります。

こうした点から、dynamic batching を利用する場合、モデルのエクスポート時に可変サイズに対応するよう設定することが推奨されます。一方で、あまりバッチ サイズを気にしなくてよいこともあります。それはどういう場合でしょうか。

バックエンドごとの違い

ここまでの内容で、モデル エクスポート時に可変バッチ サイズに対応するよう設定する必要のある条件として以下の条件があることがわかります。

• dims の先頭にバッチ サイズとして -1 を設定するケース
• max_batch_size に 1 以上の値を設定するケース
• preferred_batch_size を設定するケース

いずれのパターンも、任意のサイズのバッチを扱う必要のある状況です。一方で、バックエンド、すなわちフレームワークによってはこの制限が緩和される場合があります。ここではモデル エクスポート時のバッチ サイズの扱いと、各種設定との関係を議論します。今回は対象とする主要なバックエンドとして、以下の 4 つを取り上げます。

• TensorFlow (ここでは SavedModel を対象とします)
• PyTorch
• ONNX Runtime
• TensorRT

はじめに、おおまかに状況を整理すると以下のようになります。なお、いずれも max_batch_size の指定は dims と整合している必要があります。(たとえば、max_batch_size を指定しない場合、dims の先頭が -1 にならなくてはダメ、など)

上記の表からもわかる通り、SavedModel および ONNX が最も厳格です。バッチ サイズを固定したモデルに対して以下の設定を行う場合、モデル ロード自体に失敗します。

• dims の先頭に -1 を指定 (e.g., dims: [-1, 224, 224, 3])
• dims の先頭がモデル エクスポート時のバッチ サイズと一致しない (e.g., バッチ サイズ 1 向けのモデルに対して、dims: [4, 224, 224, 3])
• max_batch_size に 1 以上の値を指定

したがってこれら 2 つのバックエンドを利用する場合、モデルの定義と設定の整合性には特に注意する必要があります。また冒頭で述べた、自動生成される設定を利用する場合、固定バッチ サイズのモデルではうまく dynamic batching が動作しない (=クライアント側でバッチ サイズを気にする必要が生じて嬉しくない) ためか、デフォルトで dynamic batching が有効化される動作は無効になります。

次に厳格なのが TensorRT です。SavedModel や ONNX と異なり、固定バッチ サイズのモデルでも dims の先頭要素に -1 を指定したり、max_batch_size に 1 以上の値を指定できる点があります。ただしこの場合も、クライアント側のリクエストのバッチ サイズは、モデルが想定しているバッチ サイズと一致する必要がある点や、max_batch_size はモデルのバッチ サイズと一致させる必要がある点に注意が必要です。また、dims の先頭を -1 にしてモデルを可変バッチ サイズに対応させる場合にも若干注意が必要です。TensorRT で可変バッチ サイズに対応するためには、Dynamic Shapes という機能を使用するのですが、これは一定範囲内のバッチ サイズに対応する機能のため、無制限に任意のサイズを与えることはできません。たとえば、対応可能な最小バッチ サイズを 1、最適化ターゲットを 16、最大バッチ サイズを 32、というように設定します。すると、この最大サイズとして設定された 32 を超える大きなバッチを与えようとするとエラーが発生してしまうため、注意が必要となります。

最後に最も寛容なのが PyTorch で、厳密には TorchScript を利用する場合、最も制限が緩くなります。モデル エクスポートにどういった example_inputs を指定したかによらず、dims の先頭に任意のサイズを指定可能ですし、max_batch_size も任意の値を指定可能です。もちろん、dims の先頭に 1 以上の値を指定した場合はクライアント側のバッチ サイズもそれに限定されますが、それ以外の場合は基本的に任意のサイズのバッチを利用可能です。

まとめ

Triton Inference Server における、バッチ サイズの扱いについて細かく見てきました。設定可能な項目が多いことで混乱しやすい点ではありますが、モデルをエクスポートする際の設定と、max_batch_size をどう指定するかがひとつのカギになるかと思います。典型的な用途のためには、max_batch_size に想定する最大サイズを設定し、dims にはバッチ次元を除いたテンソル形状を指定、さらに dynamic_batching: {} を設定すれば OK です。もし設定自体を考えるのが大変という場合は、Triton の起動時オプションである --strict-model-config=false をご活用ください。本記事が、今後の Triton 活用の助けになれば幸いです。