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

时间:2018-11-17 02:55:56

标签: docker tensorflow tensorflow-serving custom-operator

让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

我收到以下有关找不到我的opp的错误。

  

tensorflow_serving / util / retrier.cc:37]正在加载servable:{名称:   mymodel版本:1}失败:找不到:操作类型未注册'MyOpp'   在运行在c37a4ef2d4b4上的二进制文件中。

4 个答案:

答案 0 :(得分:2)

以下是我想对我的操作进行的操作:     -生成python包装器     -也添加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部门
    • foo_op_lib添加到contrib_ops_op_lib部门
    • foo添加到contrib_py部门
  • tensorflow/tools/pip_package/BUILD
    • 将我的python目标添加到COMMON_PIP_DEPS
  • tensorflow/core/BUILD
    • 将我的内核添加到all_kernels_statically_linked。我可能对此一无所知,但它奏效了。

以下是tensorflow服务的bazel文件:

  • WORKSPACE
    • org_tensorflow更改为指向我的张量流的local_repository,而不是Google的tensorflow_http_archive

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

答案 1 :(得分:0)

您是否在要调用的BUILD文件中添加了操作库?

答案 2 :(得分:0)

您还可以使用tensorflow作为子模块或local_repository来为您的操作使用存储库中的自定义宏。

答案 3 :(得分:0)

以下是描述该操作方法的文档:https://www.tensorflow.org/tfx/serving/custom_op

最重要的是,您需要在链接有op的情况下重建tensorflow_model_server。 tensorflow_serving / model_servers / BUILD:

SUPPORTED_TENSORFLOW_OPS = [
    ...
    "//tensorflow_serving/.../...your_op"
]