DeepLabCut GPU训练报错cudnn PoolForward launch failed

核心报错信息 #

tensorflow.python.framework.errors_impl.InternalError: 2 root error(s) found.
  (0) Internal: cudnn PoolForward launch failed
         [[{{node resnet_v1_50/pool1/MaxPool}}]]
         [[GroupCrossDeviceControlEdges_0/huber_loss/num_present/broadcast_weights/assert_broadcastable/AssertGuard/Assert/data_7/_602]]
  (1) Internal: cudnn PoolForward launch failed
         [[{{node resnet_v1_50/pool1/MaxPool}}]]
0 successful operations.
0 derived errors ignored.

常见触发原因 #

• 显存不足或被异常占用：该报错触发点为ResNet50网络第一层池化计算，属于模型加载到GPU后的首批前向操作，若此时GPU剩余显存不足以支撑算子内存分配，就会抛出launch失败错误，和版本匹配度无直接关联
• cuDNN自动调优逻辑兼容问题：pip源安装的TensorFlow默认开启cuDNN算子自动寻优，针对部分算力较低的N卡（消费级老卡、笔记本移动端显卡），预编译算子的适配逻辑存在bug，会直接触发池化算子调用失败
• 多GPU跨设备通信异常：报错栈中出现GroupCrossDeviceControlEdges节点，说明多卡训练时的权重广播逻辑触发了跨设备池化算子调用，部分消费级显卡不支持P2P显存访问时会直接抛错
• GPU运行模式限制：笔记本设备开启GPU节能模式时，核心频率动态波动会导致算子提交超时，触发同类报错

按优先级执行的解决方案 #

• 第一步：排查显存占用
  训练前执行nvidia-smi查看目标GPU的空闲显存，ResNet50骨干训练batch size=8至少需要4G以上空闲显存。关闭所有占用显存的进程（开了硬件加速的浏览器、其他深度学习任务、视频剪辑软件等），如果显存余量不足，先将训练配置的batch_size调为1，关闭多卡训练选项，单卡测试能否正常启动。
• 第二步：关闭cuDNN自动调优（最高效解决方法）
  绝大多数版本匹配但仍报错的场景都是这个原因导致的。在你的DeepLabCut训练启动脚本最开头加入以下代码，禁用cuDNN自动算子选择：

  import os
  os.environ["TF_CUDNN_USE_AUTOTUNE"] = "0"
  

  保存后重新启动训练即可。
• 第三步：校验TensorFlow GPU链路有效性
  不要仅靠版本号判断环境正常，进入Python交互环境执行以下测试代码：

  import tensorflow as tf
  print(tf.config.list_physical_devices("GPU"))
  # 执行单层池化操作测试
  test_input = tf.random.normal((1, 224, 224, 3))
  test_output = tf.keras.layers.MaxPool2D()(test_input)
  print("MaxPool operator executed successfully")
  

  如果这一步复现相同报错，说明pip安装的TensorFlow自带的cuDNN依赖和系统安装的cuDNN库冲突，卸载当前TensorFlow后安装和本地CUDA版本严格对应的tensorflow-gpu包即可，不要使用CPU/GPU混装的默认发行版。
• 第四步：系统与硬件配置排查
  • 笔记本用户将显卡工作模式切换为「高性能」，关闭GPU动态节能降频选项
  • 多卡设备训练前设置环境变量CUDA_VISIBLE_DEVICES=0（将0替换为你要使用的显卡序号），强制单卡运行绕过多卡P2P通信逻辑
  • 确认安装的GPU驱动为对应CUDA版本的正式稳定版，不要使用测试版驱动

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