TensorRT中禁用cuDNN加载的技术方案解析

背景介绍

在基于TensorRT进行深度学习推理时，系统默认会加载cuDNN库以优化计算性能。然而在某些特定场景下，开发者可能需要避免cuDNN的自动加载行为，例如：

• 需要减少运行时内存占用
• 特定硬件环境下cuDNN可能引发兼容性问题
• 某些算法在原生实现下性能更优

问题现象

当使用TensorRT 8.5.1.2版本在Jetson NX平台（CUDA 11.4 + cuDNN 8.6）上执行基于C++的物体检测推理时，系统日志会显示以下信息：

[V] Trying to load shared library libcudnn.so.8
[V] Loaded shared library libcudnn.so.8
[V] Using cuDNN as plugin tactic source
[V] Using cuDNN as core library tactic source
[I] [MemUsageChange] Init cuDNN: CPU +617, GPU +586, now: CPU 927, GPU 4940 (MiB)

解决方案

方法一：通过构建配置禁用（推荐）

最彻底的解决方案是在构建TensorRT时禁用cuDNN支持。这需要在CMake配置阶段添加以下参数：

-DTRT_DISABLE_CUDNN=ON

方法二：运行时禁用策略源（适用已有二进制）

对于已编译的TensorRT二进制文件，可以通过设置策略源(tactic sources)来禁用cuDNN：

1. 使用trtexec工具： 在执行命令中添加参数：

   --tacticSources=-CUDNN
   

2. 编程接口设置： 在C++代码中可以通过以下方式配置：

   config->setTacticSources(TacticSources::kCUBLAS | TacticSources::kCUBLAS_LT);
   

技术原理

TensorRT的优化策略源(tactic sources)决定了引擎使用哪些后端库进行计算优化。默认情况下包含：

• cuDNN：针对深度神经网络优化的计算库
• cuBLAS：基础线性代数计算库
• cuBLAS-LT：轻量级cuBLAS实现

禁用cuDNN后，TensorRT将仅使用cuBLAS等基础计算库，虽然可能损失部分优化机会，但能显著减少内存占用。

注意事项

1. 性能影响：禁用cuDNN可能导致某些层的计算性能下降
2. 功能限制：部分特殊算子可能无法在没有cuDNN支持下运行
3. 内存权衡：虽然减少了cuDNN的内存占用，但可能需要更多显存来存储中间结果

替代方案

如果完全禁用cuDNN不可行，还可以考虑：

1. 使用更轻量级的cuDNN版本
2. 延迟加载cuDNN（仅在需要时加载）
3. 针对特定网络层选择性禁用cuDNN优化

通过合理配置TensorRT的策略源，开发者可以在性能和资源消耗之间取得平衡，满足不同应用场景的需求。