使用自定义操作提供 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 二进制文件构建 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]]}'

此页面包含用于向模型服务器发送 REST 请求的更完整的 API。