Creating a Pluging layer in TensorRT

AI & Data Science Deep Learning (Training & Inference) TensorRT
1

Hello,
I want to create an ArgMax layer plugin. I read the trt samples, but I don’t know how to do that! Do you have any other tutorial or example about creating a plugin layer in trt?

Thanks.

2

Hello,

Please reference the following examples for extending TensorRT functionalities by implementing custom layers using the IPluginV2 class for the C++ and Python API.

(c++) https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#example1_add_custlay_c

(python) https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#add_custom_layer_python