使用TensorFlow Hub加载Yamnet模型报错：CUDNN_STATUS_EXECUTION_FAILED

问题场景 #

使用TensorFlow 2结合tensorflow_hub加载Yamnet模型，代码如下：

yamnet_model = hub.load('https://tfhub.dev/google/yamnet/1')
wav_data, sample_rate = librosa.load(wav_file_name, sr=16000)

scores, spectogram = yamnet_model(wav_data)

环境配置：CUDA 11.4，GPU为NVIDIA GeForce MX350（计算能力6.1），cuDNN版本8.2.4。运行最后一行代码时触发NotFoundError，错误核心为无可用卷积算法，具体日志如下：

NotFoundError
图执行错误：
2023-02-28 09:56:18.617965: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1532] 已创建设备 /job:localhost/replica:0/task:0/device:GPU:0，显存1603 MB：-> 设备: 0，名称: NVIDIA GeForce MX350，PCI总线ID: 0000:2b:00.0，计算能力: 6.1
2023-02-28 09:56:34.527922: I tensorflow/stream_executor/cuda/cuda_dnn.cc:384] 已加载cuDNN版本8204
2023-02-28 09:56:34.761531: W tensorflow/core/framework/op_kernel.cc:1745] OP_REQUIRES在conv_ops.cc:1130处失败：NOT_FOUND: 无可用算法！错误信息：
  <br>CUDNN引擎1分析失败：UNKNOWN: CUDNN_STATUS_EXECUTION_FAILED<br>
in tensorflow/stream_executor/cuda/cuda_dnn.cc(4023): 'cudnnConvolutionForward( cudnn.handle(), alpha, input_nd_.handle(), input_data.opaque(), filter_.handle(), filter_data.opaque(), conv_.handle(), ToConvForwardAlgo(algo), scratch_memory.opaque(), scratch_memory.size(), beta, output_nd_.handle(), output_data.opaque())'
  <br>CUDNN引擎0分析失败：UNKNOWN: CUDNN_STATUS_EXECUTION_FAILED<br>
in tensorflow/stream_executor/cuda/cuda_dnn.cc(4023): 'cudnnConvolutionForward( cudnn.handle(), alpha, input_nd_.handle(), input_data.opaque(), filter_.handle(), filter_data.opaque(), conv_.handle(), ToConvForwardAlgo(algo), scratch_memory.opaque(), scratch_memory.size(), beta, output_nd_.handle(), output_data.opaque())'
  <br>CUDNN引擎2分析失败：UNKNOWN: CUDNN_STATUS_EXECUTION_FAILED<br>
in tensorflow/stream_executor/cuda/cuda_dnn.cc(4023): 'cudnnConvolutionForward( cudnn.handle(), alpha, input_nd_.handle(), input_data.opaque(), filter_.handle(), filter_data.opaque(), conv_.handle(), ToConvForwardAlgo(algo), scratch_memory.opaque(), scratch_memory.size(), beta, output_nd_.handle(), output_data.opaque())'
  <br>CUDNN引擎5分析失败：UNKNOWN: CUDNN_STATUS_EXECUTION_FAILED<br>
in tensorflow/stream_executor/cuda/cuda_dnn.cc(4023): 'cudnnConvolutionForward( cudnn.handle(), alpha, input_nd_.handle(), input_data.opaque(), filter_.handle(), filter_data.opaque(), conv_.handle(), ToConvForwardAlgo(algo), scratch_memory.opaque(), scratch_memory.size(), beta, output_nd_.handle(), output_data.opaque())'

解决方案 #

• 强制切换到CPU运行
  MX350是入门级移动GPU，显存仅1.6GB，无法满足模型GPU运行需求，可直接禁用GPU：

import tensorflow as tf
tf.config.set_visible_devices([], 'GPU')

yamnet_model = hub.load('https://tfhub.dev/google/yamnet/1')
wav_data, sample_rate = librosa.load(wav_file_name, sr=16000)
scores, spectogram = yamnet_model(wav_data)

• 设置GPU动态显存分配
  若想尝试GPU运行，可让TensorFlow按需分配显存，避免一次性占满：

import tensorflow as tf
gpus = tf.config.list_physical_devices('GPU')
if gpus:
    try:
        for gpu in gpus:
            tf.config.experimental.set_memory_growth(gpu, True)
    except RuntimeError as e:
        print(e)

yamnet_model = hub.load('https://tfhub.dev/google/yamnet/1')
wav_data, sample_rate = librosa.load(wav_file_name, sr=16000)
scores, spectogram = yamnet_model(wav_data)

• 检查版本兼容性
  TensorFlow 2.x对CUDA、cuDNN版本要求严格：CUDA 11.4对应TensorFlow 2.8.x/2.9.x，cuDNN 8.2.4需匹配对应TensorFlow版本。版本不兼容会导致卷积算法加载失败，建议确认后重新安装对应版本。

• 验证输入音频格式
  Yamnet要求输入为单通道16kHz采样率的float32音频数据，检查并调整输入：

# 检查输入格式
print(wav_data.shape)  # 应为一维数组
print(wav_data.dtype)  # 应为float32

# 转换格式
wav_data = librosa.to_mono(wav_data)
wav_data = wav_data.astype('float32')

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