自定义操作#

允许自定义操作的建模框架在 Triton 推理服务器中得到部分支持。自定义操作可以在构建时或启动时添加到 Triton,并对所有加载的模型可用。

TensorRT#

TensorRT 允许用户创建 自定义层,这些自定义层随后可以在 TensorRT 模型中使用。为了使这些模型在 Triton 中运行,必须使自定义层可用。

为了使自定义层对 Triton 可用,TensorRT 自定义层实现必须编译成一个或多个共享库,然后必须使用 LD_PRELOAD 加载到 Triton 中。例如,假设您的 TensorRT 自定义层被编译成 libtrtcustom.so,使用以下命令启动 Triton 可以使这些自定义层对所有 TensorRT 模型可用。

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

这种方法的局限性在于,自定义层必须与模型仓库本身分开管理。更严重的是,如果多个共享库之间存在自定义层名称冲突,则目前无法处理。

构建自定义层共享库时,重要的是使用与 Triton 中使用的 TensorRT 版本相同的版本。您可以在 Triton 发行说明 中找到 TensorRT 版本。确保您使用正确版本的 TensorRT 的一个简单方法是使用与 Triton 容器对应的 NGC TensorRT 容器。例如,如果您使用的是 24.12 版本的 Triton,请使用 24.12 版本的 TensorRT 容器。

TensorFlow#

TensorFlow 允许用户 添加自定义操作,这些自定义操作随后可以在 TensorFlow 模型中使用。您可以通过两种方式将自定义 TensorFlow 操作加载到 Triton 中

  • 在模型加载时,通过在模型配置中列出它们。

  • 在服务器启动时,通过使用 LD_PRELOAD。

要通过模型配置注册您的自定义操作库,您可以将其作为附加字段包含在内。请参阅以下配置作为示例。

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

请注意,即使模型是在运行时加载的,多个模型也可以使用自定义运算符。目前没有办法取消分配自定义运算符,因此它们将一直可用,直到 Triton 关闭。

您还可以通过 LD_PRELOAD 注册您的自定义操作库。例如,假设您的 TensorFlow 自定义操作被编译成 libtfcustom.so,使用以下命令启动 Triton 可以使这些操作对所有 TensorFlow 模型可用。

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

使用这种方法,所有 TensorFlow 自定义操作都依赖于一个 TensorFlow 共享库,该共享库在加载时必须对自定义共享库可用。在实践中,这意味着您必须确保在发出上述命令之前,/opt/tritonserver/backends/tensorflow1 或 /opt/tritonserver/backends/tensorflow2 在库路径中。有几种方法可以控制库路径,常用的一种是使用 LD_LIBRARY_PATH。您可以在 “docker run” 命令或容器内设置 LD_LIBRARY_PATH。

$ export LD_LIBRARY_PATH=/opt/tritonserver/backends/tensorflow1:$LD_LIBRARY_PATH

这种方法的局限性在于,自定义操作必须与模型仓库本身分开管理。更严重的是,如果多个共享库之间存在自定义层名称冲突,则目前无法处理。

构建自定义操作共享库时,重要的是使用与 Triton 中使用的 TensorFlow 版本相同的版本。您可以在 Triton 发行说明 中找到 TensorFlow 版本。确保您使用正确版本的 TensorFlow 的一个简单方法是使用与 Triton 容器对应的 NGC TensorFlow 容器。例如,如果您使用的是 24.12 版本的 Triton,请使用 24.12 版本的 TensorFlow 容器。

PyTorch#

Torchscript 允许用户 添加自定义操作,这些自定义操作随后可以在 Torchscript 模型中使用。通过使用 LD_PRELOAD,您可以将自定义 C++ 操作加载到 Triton 中。例如,如果您按照 pytorch/extension-script 存储库中的说明进行操作,并且您的 Torchscript 自定义操作被编译成 libpytcustom.so,则使用以下命令启动 Triton 可以使这些操作对所有 PyTorch 模型可用。由于所有 Pytorch 自定义操作都依赖于一个或多个 PyTorch 共享库,这些共享库在加载时必须对自定义共享库可用。在实践中,这意味着您必须确保在启动服务器时 /opt/tritonserver/backends/pytorch 在库路径中。有几种方法可以控制库路径,常用的一种是使用 LD_LIBRARY_PATH。

$ LD_LIBRARY_PATH=/opt/tritonserver/backends/pytorch:$LD_LIBRARY_PATH LD_PRELOAD=libpytcustom.so:${LD_PRELOAD} tritonserver --model-repository=/tmp/models ...

这种方法的局限性在于,自定义操作必须与模型仓库本身分开管理。更严重的是,如果多个共享库之间存在自定义层名称冲突,或者用于在 PyTorch 中注册它们的句柄存在冲突,则目前无法处理。

从 Triton 20.07 版本开始,TorchVision 操作 将包含在 PyTorch 后端中,因此它们不必显式添加为自定义操作。

构建自定义操作共享库时,重要的是使用与 Triton 中使用的 PyTorch 版本相同的版本。您可以在 Triton 发行说明 中找到 PyTorch 版本。确保您使用正确版本的 PyTorch 的一个简单方法是使用与 Triton 容器对应的 NGC PyTorch 容器。例如,如果您使用的是 24.12 版本的 Triton,请使用 24.12 版本的 PyTorch 容器。

ONNX#

ONNX Runtime 允许用户 添加自定义操作,这些自定义操作随后可以在 ONNX 模型中使用。要注册您的自定义操作库,您需要将其作为附加字段包含在模型配置中。例如,如果您按照 此示例 来自 microsoft/onnxruntime 存储库,并且您的 ONNXRuntime 自定义操作被编译成 libonnxcustom.so,则将以下内容添加到模型的模型配置中,可以使这些操作对该特定 ONNX 模型可用。

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

构建自定义操作共享库时,重要的是使用与 Triton 中使用的 ONNXRuntime 版本相同的版本。您可以在 Triton 发行说明 中找到 ONNXRuntime 版本。