TensorFlow 2.3.1未调用8卡GPU训练问题求助

我之前碰到过和你几乎一模一样的情况——本地有GPU但TensorFlow只认XLA_GPU，训练全走CPU，后来排查下来是环境配置和依赖版本的问题，给你几个具体的排查和解决方向：

• 检查cuDNN版本是否匹配
  TensorFlow 2.3.1对CUDA和cuDNN的版本要求很严格，虽然你CUDA 10.2是符合要求的，但还需要搭配cuDNN 7.6.5（注意不是更高版本）。如果你没装cuDNN或者版本不对，TensorFlow会跳过物理GPU，只加载XLA相关设备。你可以去NVIDIA官网下载对应版本的cuDNN，解压后把对应的文件复制到CUDA的安装目录下（比如Linux下复制到/usr/local/cuda-10.2/lib64和include文件夹）。

• 确认环境变量配置正确
  很多时候问题出在环境变量没设对，导致TensorFlow找不到CUDA库：

  • Linux/macOS下，在终端执行以下命令（可以加到~/.bashrc或~/.zshrc里永久生效）：

    export CUDA_HOME=/usr/local/cuda-10.2
    export LD_LIBRARY_PATH=$CUDA_HOME/lib64:$CUDA_HOME/extras/CUPTI/lib64:$LD_LIBRARY_PATH
    export PATH=$CUDA_HOME/bin:$PATH
    

  • Windows下，需要在系统环境变量里添加：
    • CUDA_PATH = C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.2
    • 把C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.2\bin和C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.2\libnvvp加到Path里

• 验证TensorFlow的GPU加载日志
  在代码开头加上这两行，运行后看控制台输出的设备分配日志，能找到具体的错误原因：

  import tensorflow as tf
  tf.debugging.set_log_device_placement(True)
  

  比如如果是cuDNN版本不兼容，会直接在日志里报错；如果是权限问题，也会有对应的提示。另外可以用tf.config.experimental.list_physical_devices('GPU')代替原来的命令，这个会更精准地列出物理GPU设备。

• 重新安装TensorFlow-GPU（优先用conda）
  如果你是用pip安装的，可能存在依赖冲突的问题。建议用conda重新安装，conda会自动帮你匹配所有兼容的依赖（包括CUDA、cuDNN）：

  conda create -n tf23_gpu python=3.8
  conda activate tf23_gpu
  conda install tensorflow-gpu==2.3.1
  

  这样能避免手动配置依赖的麻烦。

• 检查是否有GPU驱动或系统权限问题
  虽然你nvidia-smi能正常显示，但有时候驱动安装不完整或者当前用户没有GPU访问权限也会导致这个问题。可以试试重启机器，或者用sudo运行你的训练脚本（Linux下），看是否能识别GPU。

另外，你提到Colab能正常运行，说明你的模型代码本身没有问题，完全是本地环境的配置问题，按照上面的步骤排查应该能解决。

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