やっはろー！
某アニメの続編が出ないかなー？と思いつつ、記事を書いてます。

今回は、前回記載した「TensorRT Inference Serverの話」 から継続して、TRTIS（TensorRT Inference Server）の話を書いていきます。

目指せ、TRTIS芸人！（すいません、まだ初級者くらいです・・涙）

何について書いてるの？

TRTISの概要的な話は前回記載しましたので、今回は特定の内容にフォーカスして、記載をしていきます。

（TRTISの）CustomInstance上で、dlib(C++向けのDNNライブラリ)を利用して、顔検知・認証の処理を動かしてみました。
今回は上記処理の内容について、少しまとめてみます。

コードは？

下記にあがってます。
https://github.com/kurama554101/sample_trtis
※本記事を書いている時点では、コミットハッシュが9f1456c129fd3def0ca727f7692528cdaf639a94の状態です。

一応、TRTISを利用したサンプルコード群を上記にあげていますが、
顔検知・認証周りについては別リポジトリに分離して、専用のサービスを作っていくかもです。

構成について

TRTISを利用する場合、以下の構成を取る必要があります。
※CustomInstanceを作る場合、前述したInstanceをビルドするためのDockerコンテナが必要になるためです。

私のリポジトリでは上記コンテナをまとめて作るように、Docker-Composeを利用して、一括で環境構築をするように対応しています。

顔認証周りのコードについて

サーバー側のコードは基本的に仕組みとして利用されるため、何もいじっていません。
※Dockerコンテナに入れるライブラリとかは追加したため、新たにDockerFileは作っています。

メインとなるのは、下記です。

• 顔検知・認証処理のモデル実装（CustomInstanceの利用）
• 顔検知・顔認証処理を試すためのクライアント実装

それぞれの処理について、後述していきます。

顔検知・認証処理のモデル実装

TRTISで独自モデルを実装をする場合、以下の流れとなります。
※今回はCustomInstanceにフォーカスしたやり方を記載しています。

• CustomInstanceを継承したモデルコードを実装
• モデルをビルドするためのCMakeLists.txtなどを作成
• モデルをビルド（cmakeを利用）
• 作成したモデルの定義ファイル（config.pbtxt）の作成
• モデル（soファイルやconfig.pbtxt）をサーバー側のmodel repositoryに格納
• サーバーを起動

今回は太字にしたところを少し補足します。

CustomInstanceを継承したモデルコードを実装

顔検知・認証処理の場合、下記のコードが該当します。

• face_recognition_model.cc
• face_wrapper.cc
  • dlibの処理をラップしているクラス

今回は実装をショートカットする都合で、dlibを利用しています。
※他のDNNフレームワークやOpenCVを利用して、顔検知・認証処理を実装するのもありだと思います。
※ただし、dlib vs OpenCVでは、dlibの方がやや精度が高そうな結果も出ていたので、dlibを採用した、という背景もあります。（こちらの記事を参考にしました）

CustomInstanceの実装しないといけないAPIは下記です。

• Init
  • モデルの初期化時に一度だけ呼ばれる（config.pbtxtの更新によるモデルのリロード処理時などのケースでも呼ばれる）
• Execute
  • 推論処理実行時に都度呼ばれる（クライアントのリクエストごとに一度実行されます）

それぞれのAPIでどのような実装をしているかは実装をみていただければと思いますが、
概要を下記に記載します。

• Init
  • face_wrapperのインスタンスを初期化（内部でdlibインスタンスを持っている）
• Execute
  • リクエストがあったIDを保持する対応（ステートフル対応のため）
  • 顔検知・認証処理
  • 検知・認証結果をjsonフォーマットでレスポンスする対応

jsonフォーマットにしたのは、（1枚の画像ごとの）出力が可変長となってしまうため、
文字列としてレスポンスを返した方が実装しやすかったためとなります。
※可変長での配列を返すうまいやり方があれば、そちらを利用する方針もありかと思います。

モデルをビルドするためのCMakeLists.txtなどを作成

• CMakeLists.txt

上記のように作れば良いです。
※dlibやOpenCVなど、利用するライブラリのリンク設定や、コンパイラオプションを設定する必要があります。

ひとつ注意なのが、CUDAでdlibを動作する場合、コンパイル時の警告を対応していない箇所があるのですが、
TRTISではWerrorが設定されているため、警告がエラーとなります。

そのため、上記のCMakeLists.txtでは、Werrorを無効化する対応を取っています。
※CMAKE_CXX_FLAGSを上書きしてます。

作成したモデルの定義ファイル（config.pbtxt）の作成

• config.pbtxt

config.pbtxtはモデルの種類ごとに必要となります。

今回のface_modelでは、今後ステートフル対応をみこして、
SequentialBatcherを利用するようなモデル設計にしています。

下記のような、control_inputの部分が、SequentialBatcherを使うための設定となります。

  control_input [
    {
      name: "START"
      control [
        {
          kind: CONTROL_SEQUENCE_START
          int32_false_true: [ 0, 1 ]
        }
      ]
    },
    {
      name: "END"
      control [
        {
          kind: CONTROL_SEQUENCE_END
          int32_false_true: [ 0, 1 ]
        }
      ]
    },
    {
      name: "READY"
      control [
        {
          kind: CONTROL_SEQUENCE_READY
          int32_false_true: [ 0, 1 ]
        }
      ]
    },
    {
      name: "CORRID"
      control [
        {
          kind: CONTROL_SEQUENCE_CORRID
          data_type: TYPE_UINT64
        }
      ]
    }
  ]

SequentialBatcherを利用する場合、クライアントからのリクエストは以下のような流れになります。
* セッション開始前にstart tokenを送る
* start後に、データを送信する
* 処理完了後にend tokenを送る

例えば、start tokenを送る場合、start token部分を1に設定して、リクエストを送れば良いです。
※ここらへんのリクエスト方法はTRTISのclient SDKを利用すれば、あまり意識する必要はないです。

顔検知・顔認証処理を試すためのクライアント実装

クライアントのコードはこちらのコードとなります。

リクエストを送信時には下記のsend_requestを実行してますが、リクエスト開始時にはflagに対して、InferRequestHeader.FLAG_SEQUENCE_STARTを設定して、送信します。

• send_requestの実装

def send_request(ctx, corr_id, array, batch_size=1, start_of_sequence=False, end_of_sequence=False):
    flags = InferRequestHeader.FLAG_NONE
    if start_of_sequence:
        flags = flags | InferRequestHeader.FLAG_SEQUENCE_START
    if end_of_sequence:
        flags = flags | InferRequestHeader.FLAG_SEQUENCE_END

    result = ctx.run({ 'INPUT' : array},
                     { 'OUTPUT' : InferContext.ResultFormat.RAW },
                     batch_size=batch_size, flags=flags, corr_id=corr_id)
    return result

• send_requestの呼び出し時

    for corr_id, img_list in datas_dict.items():
        # send start request
        send_request(infer_ctx, corr_id, [np.zeros(shape, dtype=dtype)], batch_size=batch_size, start_of_sequence=True)

        # send the image data
        for img in img_list:
            send_request(infer_ctx, corr_id, [img], batch_size=batch_size) 

        # send end request
        result = send_request(infer_ctx, corr_id, [np.zeros(shape, dtype=dtype)], batch_size=batch_size, end_of_sequence=True)

        # postprocess result
        result_map[corr_id] = json.loads(result["OUTPUT"][0][0])

また、クライアント側はstreamlitを利用して、ブラウザ経由でアクセスできるような実装にしています。

動かして見る場合

顔検知・認証処理を動かす場合、動かし方はREADME.mdに記載しています。（拙い英語ですいません・・涙）

まずは環境構築

下記のようにbashを実行すれば良いです。
※初回は、TRTISのサーバー側ビルドは恐ろしく時間がかかるので、1日くらい放置しないといけないです・・。
※現状は、サーバー側のコンテナはDockerHubからpullするだけで良くなっているので、スクリプトを直すようにします。

$ bash setup_trtis_docker_containers.sh

上記スクリプトにより、1台のマシンに、3つのコンテナ（サーバー、custom backendのビルド用、クライアント）が作成されます。

ただし、サーバーとクライアントを別マシンに入れたいケースも多々あると思います。
その場合、下記のようにコマンドを実行します。

• サーバー + custom backendのビルド用コンテナを構築

$ bash setup_trtis_docker_containers.sh --only-server

• クライアントコンテナを構築

$ bash setup_trtis_docker_containers.sh --only-client

顔認証処理を動かしてみる

クライアントのコンテナに入ります。

$ docker exec -it trtis-client-container bash

さらに、顔検知・認証用のクライアントを実行させます。（クライアントのDockerコンテナ上で実行）

$ streamlit run custom_client/python/face_recognition_client.py

クライアントを実行しているマシン上で、http://localhost:8501 にアクセスすれば、ブラウザ上で処理実行が可能です。

TIPS

TRTISで対応しているGPUアーキテクチャがcompute capabilityが6.0以上である点

今回、dlibをGPUで動作させようとして、色々とハマりましたが、
特にハマったのが上記です。

TRTISのbuild.rstにも下記記載がありました。

TRTIS_MIN_COMPUTE_CAPABILITY:
By default, the inference server supports NVIDIA GPUs with CUDA compute capability 6.0 or higher.
If all framework backends included in the inference server are built to support a lower compute capability, then TRTIS can be built to support that lower compute capability by setting -DTRTIS_MIN_COMPUTE_CAPABILITY appropriately. The setting is ignored if -DTRTIS_ENABLE_GPU=OFF.

そのため、例えば、EC2のp2.xlarge（Tesla K80 -> compute capability 3.7）だと、TRTISをGPUで動作させることは出来ません。
※厳密に言えば、Custom Instanceの内部で使うライブラリをGPU用にビルドすれば、Instance GroupをKIND_CPUにしても、ライブラリ自体はGPUで動作しそうではあります。

今回は、g4dn.xlarge（Tesla T4）を利用しました。

SequentialBatcherで対応しているバッチサイズは1のみ

こちらのコードのように、現状はSequential Batcherではバッチサイズ1のみに対応してる様です。

最後に

まだまだTRTISを使い慣れていないですが、ようやく顔検知・認証処理を動かすところまではきました。
今後は下記の様な所をさぐって、実装を追加していきたいと思っています。

• 前処理をサーバー側で行う様に対応（TRTISで既にサンプルがある：image_preprocess）
• gRPCで対応（現状はHTTPのみに対応）
• Session IDをサーバー側で自動で発行できるように対応（TRTISでそのような仕組みがあるかは不明）
• k8sを利用したスケールアウトの検討

DNNのモデルをサービスに組み込む上では、DNN Servingのフレームワークはとても魅力的と言えます。
TRTIS以外にもありますが、試しに使ってみてもらえれば幸いです。