カスタム操作でTensorFlowモデルを提供する

TensorFlow には、さまざまなハードウェア タイプ (CPU、GPU など) に合わせて微調整された ops および op カーネル (実装) の広範なライブラリがあらかじめ構築されています。これらの操作は自動的に TensorFlow Serving ModelServer バイナリにリンクされ、ユーザーによる追加作業は必要ありません。ただし、ユーザーが操作で ModelServer に明示的にリンクする必要がある使用例が 2 つあります。

  • 独自のカスタム オペレーションを作成しました (例:このガイドを使用)
  • TensorFlow に同梱されていない実装済みの op を使用しています。

演算を実装したかどうかに関係なく、カスタム演算を使用してモデルを提供するには、演算のソースにアクセスする必要があります。このガイドでは、ソースを使用してカスタム操作を提供できるようにする手順を説明します。カスタム op の実装に関するガイダンスについては、 tensorflow/custom-opリポジトリを参照してください。

前提条件: Docker がインストールされている状態で、TensorFlow Servingリポジトリのクローンを作成しており、現在の作業ディレクトリがリポジトリのルートになっています。

操作ソースをサービス提供プロジェクトにコピーします

カスタム op を使用して TensorFlow Serving を構築するには、まず op ソースをサービング プロジェクトにコピーする必要があります。この例では、上記のカスタム操作リポジトリのtensorflow_zero_outを使用します。

サービスを提供するリポジトリ内に、すべてのカスタム op を格納するcustom_opsディレクトリを作成します。この例では、 tensorflow_zero_outコードのみが存在します。

mkdir tensorflow_serving/custom_ops
cp -r <custom_ops_repo_root>/tensorflow_zero_out tensorflow_serving/custom_ops

演算用の静的ライブラリを構築する

tensorflow_zero_out のBUILD ファイルには、共有オブジェクト ファイル ( .so ) を生成するターゲットが表示されます。モデルを作成してトレーニングするために、これを Python にロードします。ただし、TensorFlow Serving はビルド時に ops を静的にリンクするため、 .aファイルが必要です。したがって、このファイルを生成するビルド ルールをtensorflow_serving/custom_ops/tensorflow_zero_out/BUILDに追加します。

cc_library(
    name = 'zero_out_ops',
    srcs = [
        "cc/kernels/zero_out_kernels.cc",
        "cc/ops/zero_out_ops.cc",
    ],
    alwayslink = 1,
    deps = [
        "@org_tensorflow//tensorflow/core:framework",
    ]
)

リンクされた op を使用して ModelServer を構築します。

カスタム op を使用するモデルを提供するには、その op がリンクされた ModelServer バイナリをビルドする必要があります。具体的には、上で作成したzero_out_opsビルド ターゲットを ModelServer のBUILDファイルに追加します。

tensorflow_serving/model_servers/BUILDを編集して、 server_libターゲットに含まれるSUPPORTED_TENSORFLOW_OPSにカスタム op ビルド ターゲットを追加します。

SUPPORTED_TENSORFLOW_OPS = [
    ...
    "//tensorflow_serving/custom_ops/tensorflow_zero_out:zero_out_ops"
]

次に、Docker 環境を使用して ModelServer を構築します。

tools/run_in_docker.sh bazel build tensorflow_serving/model_servers:tensorflow_model_server

カスタム操作を含むモデルを提供する

これで、ModelServer バイナリを実行し、このカスタム オペレーションを含むモデルの提供を開始できるようになりました。

tools/run_in_docker.sh -o "-p 8501:8501" \
bazel-bin/tensorflow_serving/model_servers/tensorflow_model_server \
--rest_api_port=8501 --model_name=<model_name> --model_base_path=<model_base_path>

推論リクエストを送信して操作を手動でテストする

これで、推論リクエストをモデル サーバーに送信して、カスタム操作をテストできるようになりました。

curl http://localhost:8501/v1/models/<model_name>:predict -X POST \
-d '{"inputs": [[1,2], [3,4]]}'

このページには、 REST リクエストをモデル サーバーに送信するためのより完全な API が含まれています。