TensorRT Custom Plugin Example
Introduction
In this blog post, I would like to demonstrate how to implement and integrate a custom plugin into TensorRT using a concrete and self-contained example.
Identity ONNX Model
1 |
# Create a neural network that consists of three identity convolutional layers. |
IdentityConv Custom Plugin Implementation
1 |
|
To perform the
IdentityConv
operation, the custom plugin class has to override the
nvinfer1::IPluginV2IOExt::enqueue
method. In our case, we simply copy the input tensor to the output tensor using
cudaMemcpyAsync
.
1 |
|
In the above implementation, some of the other caveats, such as how to specify the supported custom plugin configurations, what and how to serialize and deserialize the plugin parameters, have also been demonstrated.
IdentityConv Custom Plugin Creator Implementation
1 |
|
TensorRT allows both static and dynamic registration of custom plugins. The
REGISTER_TENSORRT_PLUGIN
macro is used to register the custom plugin creator class. In this example, we use the dynamic registration method. Therefore, the
REGISTER_TENSORRT_PLUGIN
macro is commented out.
1 |
|
In the above implementation, some of the other caveats, such as how to rely on the TensorRT ONNX parser to parse the attributes of the custom ONNX node, have also been demonstrated.
Expose IdentityConv Custom Plugin Creator to TensorRT
1 |
|
If we have multiple custom plugins, we can register them all in the
getPluginCreators
function. In this example, we only have one custom plugin.
1 |
|
Build Engine With Custom Plugin
The custom plugin can also be serialized into the TensorRT engine file via the
nvinfer1::IBuilderConfig::setPluginsToSerialize
method so that the custom plugin library is not required to be loaded during the inference time.
1 |
// Create a TensorRT engine building program that builds an engine from an ONNX |
Run Engine With Custom Plugin
1 |
|
Because the output tensor values matches the input tensor values, we have successfully verified the implementation and integration of the custom plugin.
Source Code
The source code of the example can be found in my GitHub repository “TensorRT Custom Plugin Example” .