Custom Operations

Modeling frameworks that allow custom operations are partially supported by the TensorRT Inference Server. Custom operations can be added to the server at build time or at server startup and are made available to all models loaded by the server.

TensorRT

TensorRT allows a user to create custom layers which can then be used in TensorRT models. For those models to run in the inference server the custom layers must be available to the server.

To make the custom layers available to the server, the TensorRT custom layer implementations must be compiled into one or more shared libraries which are then loaded into the inference server using LD_PRELOAD. For example, assuming your TensorRT custom layers are compiled into trtcustom.so, starting the inference server with the following command makes those custom layers available to all TensorRT models loaded into the server:

$ LD_PRELOAD=trtcustom.so trtserver --model-store=/tmp/models ...

A limitation of this approach is that the custom layers must be managed separately from the model store itself. And more seriously, if there are custom layer name conflicts across multiple shared libraries there is now way to handle it.

TensorFlow

Tensorflow allows users to add custom operations which can then be used in TensorFlow models. Currently, the only way for the inference server to support custom operations is to build them into the inference server.

To build a TensorFlow custom operation into the inference server, the source code for the custom operation must be placed in src/operations/tensorflow and the BUILD file updated.

An example operation called TRTISExampleAddSub is included in the directory in files trtis_example_addsub_op.cc and trtis_example_addsub_op.cu.cc. The build rule for the TRTISExampleAddSub operation is placed in BUILD:

tf_kernel_library(
  name = "trtis_example_addsub_op",
  srcs = ["trtis_example_addsub_op.cc"],
  gpu_srcs = ["trtis_example_addsub_op.cu.cc"],
)

The all_custom_ops entry in BUILD is updated to include trtis_example_addsub_op. When adding a new custom operation similar modifications must be made to BUILD.

After making these changes, build the server and the custom operations will be available to every TensorFlow model loaded into the server.