TensorRT项目中的GroupNormalizationPlugin插件使用问题解析

问题背景

在TensorRT项目中，用户尝试将一个包含GroupNormalization层的PyTorch模型转换为TensorRT引擎时遇到了插件加载失败的问题。具体表现为在解析ONNX模型时，系统提示无法找到"GroupNormalizationPlugin"插件。

问题现象

用户在Jetson AGX Orin设备上构建了TensorRT OSS版本，并尝试通过trtexec工具解析一个经过修改的ONNX模型。该模型原本包含InstanceNorm层，在ONNX转换过程中被替换为"GroupNormalizationPlugin"节点。然而，TensorRT解析器报告无法找到该插件。

问题分析

通过深入分析，我们发现导致该问题的根本原因有两个：

1. 插件版本格式错误：在ONNX模型中，plugin_version属性被错误地设置为整数类型(int)，而TensorRT插件注册系统期望接收的是字符串类型(string)的版本号。

2. 依赖库版本不匹配：GroupNormalizationPlugin插件依赖于libcudnn.so.8库，而系统中可能安装的是较新版本的cuDNN库(如libcudnn.so.9)。

解决方案

针对上述问题，我们提供了以下解决方案：

1. 修正插件版本格式： 在修改ONNX模型时，确保将plugin_version属性设置为字符串类型：

   attrs['plugin_version'] = "1"  # 正确：字符串类型
   # 而不是
   attrs['plugin_version'] = 1    # 错误：整数类型
   

2. 解决库依赖问题： 创建适当的符号链接并确保库路径正确：

   ln -s libcudnn.so.9 libcudnn.so.8
   export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/path/to/cudnn/library
   

技术要点

1. TensorRT插件系统：TensorRT通过插件机制支持自定义操作，插件需要正确注册并实现特定的接口才能被识别和使用。

2. 版本兼容性：在深度学习框架和库的集成过程中，版本兼容性至关重要，包括插件版本号的格式和依赖库的版本。

3. ONNX模型修改：使用工具如onnx-graphsurgeon修改ONNX模型时，需要特别注意属性值的类型和格式要求。

验证结果

通过上述修正后，用户成功加载了GroupNormalizationPlugin插件，并且验证了TensorRT引擎的输出与原始PyTorch模型和ONNX模型的输出完全匹配（bit-match）。

最佳实践建议

1. 在开发TensorRT插件时，建议添加详细的日志输出，便于调试插件加载过程。

2. 对于依赖库，建议明确文档说明所需的版本，并在构建时进行版本检查。

3. 在修改ONNX模型属性时，应参考目标插件的具体实现，确保属性类型和值符合预期。

通过本案例的分析和解决，我们不仅解决了具体的技术问题，也为类似场景下的TensorRT插件集成提供了有价值的参考经验。