TensorFlow 2.1 C++加载YOLOv3模型时触发CUDNN_STATUS_INTERNAL_ERROR

首先帮你梳理下核心问题：你遇到的CUDNN_STATUS_INTERNAL_ERROR和卷积算法获取失败，不一定是6G显存不足导致的——YOLOv3默认416x416输入尺寸下，6G显存完全足够支撑推理，我们从几个方向逐一排查解决：

1. 优先排查版本兼容性冲突 #

你的环境里有个潜在的版本不匹配点：

• 你安装的是CUDA 10.1，但nvidia-smi显示驱动对应的CUDA版本是11.0。虽然CUDA 10.1兼容部分高版本驱动，但驱动版本过高可能和cuDNN 7.6.5存在适配问题。
• 验证建议：可以运行TensorFlow官方提供的CUDA检测示例，或者尝试重新安装与CUDA 10.1严格匹配的驱动（比如418.39或440.33版本），确保驱动、CUDA、cuDNN三者完全符合TensorFlow 2.1的官方要求。

2. 优化GPU显存管理（不止开启内存增长） #

你已经设置了allow_growth=true，但可以再补充以下操作：

• 限制显存占用比例：在ConfigProto中添加显存占比设置，避免程序尝试占用超出实际可用的显存：

tensorflow::ConfigProto config;
config.mutable_gpu_options()->set_allow_growth(true);
// 限制使用80%的GPU显存，留出缓冲空间
config.mutable_gpu_options()->set_per_process_gpu_memory_fraction(0.8);

• 清理不必要的显存占用：从你的nvidia-smi输出看，桌面环境（Xorg、gnome-shell）占用了近500MiB显存。如果是在桌面环境下运行推理程序，可以尝试切换到纯命令行模式（Ctrl+Alt+F3），杀掉桌面进程后再运行，释放更多可用显存。

3. 检查YOLOv3的输入尺寸设置 #

如果你使用的是608x608甚至更大的输入尺寸，显存占用会显著提升。建议先切换到YOLOv3默认的416x416输入尺寸测试，如果能正常运行，再逐步调整到你需要的尺寸。

4. 确保环境变量配置正确 #

运行程序前，手动配置环境变量，确保CUDA和cuDNN的库路径被正确加载：

export LD_LIBRARY_PATH=/usr/local/cuda/lib64:/usr/local/cuda/include:$LD_LIBRARY_PATH
export CUDA_VISIBLE_DEVICES=0

这可以避免程序找不到cuDNN库或者误调用其他GPU设备的问题。

5. 验证cuDNN基础功能是否正常 #

可以写一个简单的TensorFlow C++程序测试cuDNN是否能正常初始化，比如执行一个基础卷积操作：

#include <tensorflow/core/public/session.h>
#include <tensorflow/core/platform/env.h>
#include <tensorflow/core/framework/tensor.h>

using namespace tensorflow;

int main() {
    Session* session;
    Status status = NewSession(SessionOptions(), &session);
    if (!status.ok()) {
        std::cout << status.ToString() << std::endl;
        return 1;
    }

    // 构建简单的卷积测试
    Tensor input(DT_FLOAT, TensorShape({1, 416, 416, 3}));
    input.flat<float>().setZero();

    auto root = tensorflow::Scope::NewRootScope();
    auto input_tensor = ops::Placeholder(root.WithOpName("input"), DT_FLOAT);
    auto conv = ops::Conv2D(root.WithOpName("conv"), input_tensor, 
                            ops::Const(root, {3,3,3,32}, DT_FLOAT), 
                            {1,1,1,1}, "SAME");

    GraphDef graph_def;
    status = root.ToGraphDef(&graph_def);
    if (!status.ok()) {
        std::cout << status.ToString() << std::endl;
        return 1;
    }

    status = session->Create(graph_def);
    if (!status.ok()) {
        std::cout << status.ToString() << std::endl;
        return 1;
    }

    std::vector<std::pair<string, Tensor>> inputs = {{"input", input}};
    std::vector<Tensor> outputs;
    status = session->Run(inputs, {"conv"}, {}, &outputs);
    if (!status.ok()) {
        std::cout << status.ToString() << std::endl;
        return 1;
    }

    std::cout << "Convolution executed successfully!" << std::endl;
    session->Close();
    delete session;
    return 0;
}

如果这个测试程序也报错，说明问题出在CUDA/cuDNN的环境配置上，而非YOLOv3模型本身。

总的来说，6G显存完全足够运行YOLOv3推理，重点排查版本兼容性和环境配置问题即可。

内容的提问来源于stack exchange，提问作者user6709725