docker - 如何在 Tensorflow 服务器上使用自定义操作

Question

让 tensorflow 模型服务器识别我的自定义操作的理想方式是什么？

我按照本指南编写了一个自定义操作: https://www.tensorflow.org/guide/extend/op

我可以通过调用tf.load_op_library来使用opp，但是当我尝试并运行tensorflow_model_server

tensorflow_model_server --port=9000 \
                        --model_name=mymodel \
                        --model_base_path=/serving/mymodel

我收到以下关于无法找到我的对手的错误。

tensorflow_serving/util/retrier.cc:37] Loading servable: {name: mymodel version: 1} failed: Not found: Op type not registered 'MyOpp' in binary running on c37a4ef2d4b4.

Answer 1

以下是我想用我的操作做的事情: - 生成Python包装器 - 在 pip 包中添加 op - 将我的操作链接到tensorflow，以便tensorflow-serving可以执行该操作

我将操作放在tensorflow/contrib/foo 中。这是源代码树的样子

.
├── BUILD
├── LICENSE
├── __init__.py
├── foo_op.cc
├── foo_op_gpu.cu.cc
└── foo_op.h

我的 __init__.py 文件导入了生成的包装器

from tensorflow.contrib.sampling.ops.gen_foo import *

我在 tensorflow/contrib/__init__.py

中添加了导入

from tensorflow.contrib import foo

这是我的 tensorflow/contrib/foo/BUILD 文件:

licenses(["notice"])  # Apache 2.0
exports_files(["LICENSE"])

package(default_visibility = ["//visibility:public"])

load("//tensorflow:tensorflow.bzl", "tf_custom_op_py_library")
load("//tensorflow:tensorflow.bzl", "tf_gen_op_libs")
load("//tensorflow:tensorflow.bzl", "tf_gen_op_wrapper_py")
load("//tensorflow:tensorflow.bzl", "tf_kernel_library")

tf_kernel_library(
    name = "foo_op_kernels",
    prefix = "foo",
    alwayslink = 1,
)
tf_gen_op_libs(
    op_lib_names = ["foo"],
)
tf_gen_op_wrapper_py(
    name = "foo",
    visibility = ["//visibility:public"],
    deps = [
        ":foo_op_kernels",
    ],
)
tf_custom_op_py_library(
    name = "foo_py",
    srcs = [
        "__init__.py",
    ],
    kernels = [
        ":foo_op_kernels",
    ],
    srcs_version = "PY2AND3",
    deps = [
        ":foo",
        "//tensorflow/contrib/util:util_py",
        "//tensorflow/python:common_shapes",
        "//tensorflow/python:framework_for_generated_wrappers",
        "//tensorflow/python:platform",
        "//tensorflow/python:util",
    ],
)

这是我必须接触才能使其正常工作的 TensorFlow bazel 文件。

• tensorflow/contrib/BUILD
  • 将 foo_op_kernels 添加到 contrib_kernels deps
  • 将 foo_op_lib 添加到 contrib_ops_op_lib deps
  • 将 foo 添加到 contrib_py deps
• tensorflow/tools/pip_package/BUILD
  • 将我的 python 目标添加到 COMMON_PIP_DEPS
• tensorflow/core/BUILD
  • 将我的内核添加到 all_kernels_statically_linked。我可能有点太过分了，但它确实有效。

以下是提供 bazel 文件的 tensorflow :

• 工作空间
  • 将org_tensorflow更改为指向我的tensorflow的local_repository，而不是Google的tensorflow_http_archive

然后我修改:tensorflow_serving/tools/docker/Dockerfile.devel-gpu来克隆我的tensorflow和tensorflow-serving版本。