Lei Mao's Log Book
·
TensorRT 自定义插件示例
💡
原文英文，约5500词，阅读约需20分钟。
📝
内容提要
本文讨论了在TensorRT中实现和集成自定义插件的方法，并提供了一个具体示例。自定义插件用于加速在NVIDIA GPU上的深度学习推理，并允许用户为TensorRT不支持的神经网络层创建自定义插件。文章提供了创建自定义插件、实现插件类以及在TensorRT中注册插件的代码示例。还演示了如何使用自定义插件构建引擎并运行推理。
🎯
关键要点
- TensorRT是一个高性能的深度学习推理SDK，支持用户为不支持的神经网络层创建自定义插件。
- 本文提供了如何在TensorRT中实现和集成自定义插件的具体示例。
- 创建了一个简单的身份ONNX模型，其中包含三个卷积节点，第二个节点被替换为自定义的IdentityConv节点。
- 自定义插件类需要从nvinfer1::IPluginV2IOExt或nvinfer1::IPluginV2DynamicExt类派生。
- IdentityConv插件的enqueue方法被重写，以实现输入张量到输出张量的简单复制。
- 自定义插件创建者类需要从nvinfer1::IPluginCreator类派生，并实现createPlugin和deserializePlugin方法。
- TensorRT允许静态和动态注册自定义插件，示例中使用动态注册方法。
- 自定义插件库可以通过nvinfer1::IPluginRegistry::loadLibrary方法动态加载。
- 自定义插件可以通过nvinfer1::IBuilderConfig::setPluginsToSerialize方法序列化到TensorRT引擎文件中。
- 在推理时，可以直接从引擎文件反序列化引擎，而无需加载自定义插件库。
➡️
