TensorRT - 使用自定义层扩展 TensorRT

使用自定义层扩展 TensorRT

NVIDIA TensorRT 支持许多层,其功能不断扩展;但是,在某些情况下,支持的层可能无法满足模型的特定需求。在这种情况下,可以通过实现自定义层(通常称为插件)来扩展 TensorRT。

TensorRT 包含可加载到应用程序中的标准插件。有关开源插件列表,请参阅 GitHub: TensorRT 插件。

要在应用程序中使用标准 TensorRT 插件,必须加载 libnvinfer_plugin.so(Windows 上为 nvinfer_plugin.dll)库,并且必须通过在应用程序代码中调用 initLibNvInferPlugins 来注册所有插件。有关这些插件的更多信息,请参阅 NvInferPlugin.h 文件。

如果这些插件不满足您的需求,您可以编写并添加自己的插件。

使用 C++ API 添加自定义层

确保 TensorRT 正确识别您的插件有四个步骤:

1. 从 TensorRT 的插件基类之一实现插件