使用自訂運算元提供 TensorFlow 模型

TensorFlow 已預先建構,其中包含廣泛的運算元和運算元核心 (實作),並針對不同的硬體類型 (CPU、GPU 等) 進行微調。這些運算會自動連結到 TensorFlow Serving ModelServer 二進位檔中,使用者無需執行額外作業。但是,有兩種使用情境需要使用者明確地將運算元連結到 ModelServer 中

  • 您已編寫自己的自訂運算元 (例如,使用本指南)
  • 您正在使用 TensorFlow 未隨附的已實作運算元

無論您是否實作運算元,為了使用自訂運算元提供模型,您都需要存取運算元的來源。本指南將逐步引導您使用來源,讓自訂運算元可用於服務。如需實作自訂運算元的指南,請參閱 tensorflow/custom-op 存放區。

先決條件:在安裝 Docker 的情況下,您已複製 TensorFlow Serving 存放區,且您目前的工作目錄是存放區的根目錄。

將運算元來源複製到 Serving 專案中

為了使用您的自訂運算元建構 TensorFlow Serving,您首先需要將運算元來源複製到您的 serving 專案中。在本範例中,您將使用上方提及的 custom-op 存放區中的 tensorflow_zero_out

在 serving 存放區中,建立 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 會在建構時靜態連結運算元,且需要 .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",
    ]
)

建構連結運算元的 ModelServer

若要提供使用自訂運算元的模型,您必須建構連結該運算元的 ModelServer 二進位檔。具體來說,您會將上方建立的 zero_out_ops 建構目標新增至 ModelServer 的 BUILD 檔案。

編輯 tensorflow_serving/model_servers/BUILD,將您的自訂運算元建構目標新增至 SUPPORTED_TENSORFLOW_OPS,此目標包含在 server_lib 目標中

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 https://127.0.0.1:8501/v1/models/<model_name>:predict -X POST \
-d '{"inputs": [[1,2], [3,4]]}'

本頁面包含更完整的 API,可用於將 REST 要求傳送到模型伺服器。