Custom Operations

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

TensorRT

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

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

$ LD_PRELOAD=libtrtcustom.so tritonserver --model-repository=/tmp/models ...

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

When building the custom layer shared library it is important to use the same version of TensorRT as is being used in Triton. You can find the TensorRT version in the Triton Release Notes. A simple way to ensure you are using the correct version of TensorRT is to use the NGC TensorRT container corresponding to the Triton container. For example, if you are using the 20.03 version of Triton, use the 20.03 version of the TensorRT container.

TensorFlow

Tensorflow allows users to add custom operations which can then be used in TensorFlow models. By using LD_PRELOAD you can load your custom TensorFlow operations into Triton. For example, assuming your TensorFlow custom operations are compiled into libtfcustom.so, starting Triton with the following command makes those operations available to all TensorFlow models:

$ LD_PRELOAD=libtfcustom.so tritonserver --model-repository=/tmp/models ...

All TensorFlow custom operations depend on a TensorFlow shared library that must be available to the custom shared library when it is loading. In practice this means that you must make sure that /opt/tritonserver/lib/tensorflow is on the library path before issuing the above command. There are several ways to control the library path and a common one is to use the LD_LIBRARY_PATH. You can set LD_LIBRARY_PATH in the “docker run” command or inside the container:

$ export LD_LIBRARY_PATH=/opt/tritonserver/lib/tensorflow:$LD_LIBRARY_PATH

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

When building the custom operations shared library it is important to use the same version of TensorFlow as is being used in Triton. You can find the TensorFlow version in the Triton Release Notes. A simple way to ensure you are using the correct version of TensorFlow is to use the NGC TensorFlow container corresponding to the Triton container. For example, if you are using the 20.03 version of Triton, use the 20.03 version of the TensorFlow container.

PyTorch

Torchscript allows users to add custom operations which can then be used in Torchscript models. By using LD_PRELOAD you can load your custom C++ operations into Triton. For example, if you follow the instructions in the pytorch/extension-script repository and your Torchscript custom operations are compiled into libpytcustom.so, starting Triton with the following command makes those operations available to all PyTorch models:

$ LD_PRELOAD=libpytcustom.so tritonserver --model-repository=/tmp/models ...

All PyTorch custom operations depend on one or more PyTorch shared libraries that must be available to the custom shared library when it is loading. In practice this means that you must make sure that /opt/tritonserver/lib/pytorch is on the library path before issuing the above command. There are several ways to control the library path and a common one is to use the LD_LIBRARY_PATH. You can set LD_LIBRARY_PATH in the “docker run” command or inside the container:

$ export LD_LIBRARY_PATH=/opt/tritonserver/lib/pytorch:$LD_LIBRARY_PATH

A limitation of this approach is that the custom operations must be managed separately from the model repository itself. And more seriously, if there are custom layer name conflicts across multiple shared libraries or the handles used to register them in PyTorch there is currently no way to handle it.

When building the custom operations shared library it is important to use the same version of PyTorch as is being used in Triton. You can find the PyTorch version in the Triton Release Notes. A simple way to ensure you are using the correct version of PyTorch is to use the NGC PyTorch container corresponding to the Triton container. For example, if you are using the 20.03 version of Triton, use the 20.03 version of the PyTorch container.

ONNX

ONNXRuntime allows users to add custom operations which can then be used in ONNX models. To register your custom operations library you need to include it in the model configuration as an additional field. For example, if you follow this example in the `microsoft/onnxruntime<https://github.com/microsoft/onnxruntime>`_ repository and your ONNXRuntime custom operations are compiled into libonnxcustom.so, adding the following to the model configuraion of your model makes those operations available to that specific ONNX model:

$ model_operations { op_library_filename: "/path/to/libonnxcustom.so" }

When building the custom operations shared library it is important to use the same version of ONNXRuntime as is being used in Triton. You can find the ONNXRuntime version in the Triton Release Notes.