TensorFlow识别GPU但训练时未利用，出现资源耗尽错误求助

可能原因及对应解决方案 #

• 显存分配策略问题
  TensorFlow默认尝试占用全部GPU显存，若显存不足会自动 fallback 到CPU，导致CPU内存耗尽。可通过以下代码设置按需分配显存：

  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)
          print(f"{len(gpus)} 物理GPU，{len(tf.config.list_logical_devices('GPU'))} 逻辑GPU")
      except RuntimeError as e:
          print(e)
  

  也可固定分配显存比例（根据GPU实际显存调整）：

  tf.config.set_logical_device_configuration(
      gpus[0],
      [tf.config.LogicalDeviceConfiguration(memory_limit=4096)] # 限制为4GB
  )
  

• 模型/数据未部署到GPU
  确保输入数据和模型都加载到GPU设备上：

  # 转移数据到GPU
  x_train = x_train.gpu()
  y_train = y_train.gpu()
  
  # 指定GPU构建并训练模型
  with tf.device('/GPU:0'):
      model = tf.keras.Sequential([
          # 你的模型层定义
      ])
      model.compile(optimizer='adam', loss='sparse_categorical_crossentropy')
      model.fit(x_train, y_train, epochs=5)
  

• CUDA/cuDNN版本不兼容
  TensorFlow 2.9.0官方推荐的CUDA版本为11.2，cuDNN为8.1.0，CUDA 11.6可能存在兼容性问题。建议：

  • 卸载当前CUDA 11.6，安装官方匹配版本
  • 重新配置环境变量（确保CUDA_PATH、PATH、LD_LIBRARY_PATH等指向正确路径）

• GPU显存不足导致自动切换CPU
  若GPU显存过小，模型或batch size过大时会触发CPU fallback：

  • 减小batch size（如从32调整为8/16）
  • 简化模型结构（减少层数、降低神经元数量）
  • 启用混合精度训练减少显存占用：

    from tensorflow.keras.mixed_precision import set_global_policy
    set_global_policy('mixed_float16')
    

• 验证GPU实际使用情况
  运行以下代码确认TensorFlow是否真正调用GPU：

  print("TensorFlow是否基于CUDA构建：", tf.test.is_built_with_cuda())
  print("GPU是否可用：", tf.test.is_gpu_available())
  
  # 强制在GPU上训练，查看报错信息
  with tf.device('/GPU:0'):
      model.fit(x_train, y_train, epochs=1)
  

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