tensorrt.plugin.register#

tensorrt.plugin.register(plugin_id: str, lazy_register: bool = False) → Callable#

Wraps a function to register and describe a TensorRT plugin’s IO characteristics. In addition, a complete plugin at least needs an trt.plugin.impl function to be registered.

This API is only intended to be used as a decorator. The decorated function must have type hints for all inputs as well as return value.

(inp0: TensorDesc, inp1: TensorDesc, ..., attr0: SupportedAttrType, attr1: SupportedAttrType, ...) -> Union[TensorDesc, Tuple[TensorDesc]]

  • Input tensors are declared first, each described by a tensor descriptor TensorDesc.

  • Plugin attributes are declared next. “SupportedAttrType” must be one of:

    • Supported built-in types: int, float, str, bool, bytes (Note: Lists/tuples of these types are not supported)

    • 1-D Numpy arrays of the following types: int8, int16, int32, int64, float16, float32, float64, bool. These must be annotated with ‘numpy.typing.NDArray[dtype]’, where ‘dtype’ is the expected numpy dtype.

  • If the plugin has only one output, the return annotation could be TensorDesc. Tuple[TensorDesc] could be used for any number of outputs.

By default, the plugin will be immediately registered in the TRT plugin registry. Use the lazy_register argument to change this.

Parameters

  • plugin_id – An ID for the plugin in the form “{namespace}::{name}”, e.g. “my_project::add_plugin”. The namespace is used to avoid collisions so using your product/project name is recommended.

  • lazy_register – During plugin development/when building engine, lazy registration may be used to delay plugin registration until the plugin is explicitly instantiated using trt.plugin.op.ns.plugin_name(…)

Registration of an elementwise plugin (output has same characteristics as the input)#
1import tensorrt.plugin as trtp
2
3@trtp.register("my::add_plugin")
4def add_plugin_desc(inp0: trtp.TensorDesc, block_size: int) -> Tuple[trtp.TensorDesc]:
5    return inp0.like()