如何强制Keras使用TensorFlow GPU后端（CPU/GPU版共存场景）

可行解决方案如下 #

• 第一步：验证TensorFlow GPU检测状态

  运行以下代码确认GPU是否被当前加载的TensorFlow识别：

  import tensorflow as tf
  # TF 1.x版本可直接运行该命令验证GPU可用性，返回True即为识别成功
  print(tf.test.is_gpu_available())
  # 启动会话时打印设备分配日志，可确认运算实际调用的设备
  sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
  

  如果返回False或者日志显示运算均调用CPU，说明GPU版本未被正确加载，继续执行后续操作。

• 第二步：设置CUDA可见设备强制调用GPU

  配置环境变量屏蔽CPU版本TensorFlow的优先级，让TensorFlow只能识别到GPU设备：

  • Linux/macOS终端运行脚本前执行：export CUDA_VISIBLE_DEVICES=0
  • Windows命令行运行脚本前执行：set CUDA_VISIBLE_DEVICES=0
    也可以直接在代码最开头添加环境变量配置：

  import os
  os.environ["CUDA_VISIBLE_DEVICES"] = "0"
  import tensorflow as tf
  

• 第三步：调整包导入优先级

  当两个版本共存时，CPU版TensorFlow的安装路径可能在系统导入路径中优先级更高，可手动将GPU版路径前置：

  1. 执行pip show tensorflow-gpu，复制输出结果中Location字段对应的路径
  2. 在代码最开头插入路径配置，确保导入时优先读取GPU版本：

  import sys
  # 引号中替换为你刚才复制的tensorflow-gpu实际安装路径
  sys.path.insert(0, "/your/env/path/lib/pythonx.x/site-packages")
  # 先导入tensorflow再导入keras，确保keras使用已加载的GPU版TensorFlow
  import tensorflow as tf
  import keras
  

• 第四步：运行时强制指定GPU设备

  可将模型运算代码块手动绑定到GPU设备运行：

  with tf.device('/gpu:0'):
      # 模型构建、训练代码均放在该代码块下
      model.fit(train_x, train_y, epochs=10)
  

以上方法均不需要删除CPU版本TensorFlow，可实现版本共存场景下强制调用GPU版本。

内容的提问来源于stack exchange，提问作者Nagabhushan S N