如何使用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]正在加载servable：｛name:mymodel版本：1｝失败：未找到：Op类型未注册"MyOpp"在c37a4ef2d4b4上运行的二进制文件中。

Answer 1

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

我把我的操作放在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_kernelsdeps
  • 将foo_op_lib添加到contrib_ops_op_libdeps
  • 将foo添加到contrib_pydeps
• tensorflow/tools/pip_package/BUILD
  • 将我的python目标添加到COMMON_PIP_DEPS
• tensorflow/core/BUILD
  • 将我的内核添加到all_kernels_statically_linked。我可能对这个太过火了，但它起了作用

以下是提供集市文件的tensorflow：

• WORKSPACE
  • 将org_tensorflow更改为指向我的tensorflow的local_repository，而不是谷歌的tensorflow_http_archive

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

Answer 2

下面是一个文档，描述了如何做到这一点：https://www.tensorflow.org/tfx/serving/custom_op

最重要的是，您需要在链接了操作的情况下重建tensorflow_model_server。tensorflow_serving/model_servers/BUILD:

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

Answer 3

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

Answer 4

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