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

Posted

技术标签:

【中文标题】如何在 tensorflow 服务器上使用自定义操作【英文标题】:How to use custom operations with tensorflow server 【发布时间】:2018-11-17 02:55:56 【问题描述】:

让 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: name: mymodel 版本:1 失败:未找到:操作类型未注册“MyOpp” 在 c37a4ef2d4b4 上运行的二进制文件中。

【问题讨论】:

作为旁注,我在 tensorflow 的自定义操作 repo github.com/tensorflow/custom-op/issues/3 上提出了类似的问题 【参考方案1】:

以下是我想用我的操作做的事情: - 生成 python 包装器 - 将 op 也添加到 pip 包中 - 将我的操作链接到 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/BUILDfoo_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 文件:

WORKSPACEorg_tensorflow 更改为指向我的张量流的local_repository 而不是google 的tensorflow_http_archive

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

【讨论】:

您对这些步骤有详细的解释吗?一个完整的例子可能会更好?谢谢【参考方案2】:

这是一个描述如何做到这一点的文档:https://www.tensorflow.org/tfx/serving/custom_op

底线是您需要重建 tensorflow_model_server 并链接您的操作。 tensorflow_serving/model_servers/BUILD:

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

【讨论】:

【参考方案3】:

你是否在 BUILD 文件中添加了你的 op lib 到你想调用它的地方?

【讨论】:

谢谢你的建议让我开始了。我不得不做更多的事情来让一切正常工作。【参考方案4】:

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

【讨论】:

yifeif,您能举个例子吗? 我还想看看这种方法的一个例子:) 听起来更干净/更容易。

以上是关于如何在 tensorflow 服务器上使用自定义操作的主要内容,如果未能解决你的问题,请参考以下文章

如何在 tensorflow 服务 (gpu) 中加载自定义动态库 (*.so)?

如何在 Tensorflow 2.x Keras 自定义层中使用多个输入?

如何使用 tensorflow-serving 发布自定义(非 tensorflow)模型?

如何在 TensorFlow 中使用 Hugging Face Transformers 库对自定义数据进行文本分类?

如何使用自定义protobuf构建tensorflow 1.13.1?

如何在Tensorflow中使用自定义/非默认tf.Graph正确的方法?