How to load custom operation library on TensorFlow Serving

Updated on 02 Dec 2018

Suppose that you define your customer operation library in file my_op.cc that you want it to be compiled in Tensorflow Serving. There are only two steps to get it done.

Fisrt step

Add one line in tensorflowserving/modelservers/BUILD as below:

SUPPORTED_TENSORFLOW_OPS = [
    "@org_tensorflow//tensorflow/contrib:contrib_kernels",
    "@org_tensorflow//tensorflow/contrib:contrib_ops_op_lib",
    "//tensorflow_serving/mzy_op:my_op.so"  # Add this line
]

Second step

Edit my_ops/BUILD

package(
    default_visibility = [
        "//tensorflow_serving:internal",
    ],
    features = ["-layering_check"],
)

cc_library(
    name = "my_op.so",
    visibility = ["//visibility:public"],
    srcs = [
            "my_op.cc",
           ] ,
    copts = ["-std=c++11"],
    deps = ["@org_tensorflow//tensorflow/core:framework_headers_lib",
            "@org_tensorflow//tensorflow/core/util/ctc",
            "@org_tensorflow//third_party/eigen3",
    ],
    alwayslink = 1,
)

In the end, you will have a structure like this:

- tensorflow_serving
    ...
    - model_servers
        ...
        - BUILD
        ...
    - my_ops
        - BUILD
        - my_op.cc
    ...