Jupyter Notebook无法识别GPU但Conda环境中可正常检测

问题场景 #

通过VSCode SSH连接远程服务器，创建了名为myspace的Anaconda环境并安装TensorFlow 2.6.0：

• 终端激活环境后，TensorFlow能正常识别GPU：

(myspace) user@server:~$ python
Python 3.6.13 |Anaconda, Inc.| (default, Jun  4 2021, 14:25:59) 
[GCC 7.5.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import tensorflow as tf
>>> print("Num GPUs:", len(tf.config.experimental.list_physical_devices('GPU')))
Num GPUs: 1

• 但在Jupyter Notebook中选择该环境作为内核时，GPU检测失败：

print("Num GPUs:", len(tf.config.experimental.list_physical_devices('GPU')))
# 输出：Num GPUs: 0

已确认CUDA Toolkit 11.2、CUDNN 8安装正常，TensorFlow编译信息显示CUDA构建有效：

tf.sysconfig.get_build_info()
# 输出：
OrderedDict([('cpu_compiler', '/usr/bin/gcc-5'),
             ('cuda_compute_capabilities',
              ['sm_35', 'sm_50', 'sm_60', 'sm_70', 'sm_75', 'compute_80']),
             ('cuda_version', '11.2'),
             ('cudnn_version', '8'),
             ('is_cuda_build', True),
             ('is_rocm_build', False),
             ('is_tensorrt_build', True)])

尝试在Jupyter中获取os.environ["LD_LIBRARY_PATH"]抛出KeyError，推测环境路径未加载，已在~/.bashrc中配置：

export PATH=/home/user/cuda-11.2/bin:${PATH}
export LD_LIBRARY_PATH=/home/user/cuda-11.2/lib64:${LD_LIBRARY_PATH}

解决方案 #

1. 让Jupyter直接加载CUDA环境变量 #

Jupyter默认可能以非交互式shell启动，不会加载~/.bashrc的配置，可直接在Jupyter配置中指定环境变量：

• 激活myspace环境，生成Jupyter配置文件：

  conda activate myspace
  jupyter notebook --generate-config
  

• 打开配置文件~/.jupyter/jupyter_notebook_config.py，找到c.NotebookApp.env项，取消注释并修改为：

  import os
  c.NotebookApp.env = {
      'LD_LIBRARY_PATH': '/home/user/cuda-11.2/lib64:' + os.environ.get('LD_LIBRARY_PATH', ''),
      'PATH': '/home/user/cuda-11.2/bin:' + os.environ.get('PATH', '')
  }
  

• 重启Jupyter Notebook后重新测试GPU检测。

2. 在Conda环境内配置专属环境变量 #

这种方式不依赖系统shell配置，环境隔离性更强：

• 激活myspace环境：

  conda activate myspace
  

• 创建环境激活时的配置目录与文件：

  mkdir -p $CONDA_PREFIX/etc/conda/activate.d
  touch $CONDA_PREFIX/etc/conda/activate.d/env_vars.sh
  

• 编辑env_vars.sh写入CUDA路径配置：

  export PATH=/home/user/cuda-11.2/bin:$PATH
  export LD_LIBRARY_PATH=/home/user/cuda-11.2/lib64:$LD_LIBRARY_PATH
  

• 创建环境退出时的清理配置：

  mkdir -p $CONDA_PREFIX/etc/conda/deactivate.d
  touch $CONDA_PREFIX/etc/conda/deactivate.d/env_vars.sh
  

• 编辑退出时的env_vars.sh移除CUDA路径：

  export PATH=$(echo $PATH | sed -e 's|/home/user/cuda-11.2/bin:||')
  export LD_LIBRARY_PATH=$(echo $LD_LIBRARY_PATH | sed -e 's|/home/user/cuda-11.2/lib64:||')
  

• 重新激活环境：conda deactivate && conda activate myspace，重启Jupyter测试。

3. 重新注册Jupyter内核，确保环境关联正确 #

有时内核显示名称正确，但实际关联的Python路径有误：

• 激活myspace环境，重装ipykernel：

  conda activate myspace
  conda install ipykernel --force-reinstall
  

• 重新注册内核：

  python -m ipykernel install --user --name myspace --display-name "Python (myspace)"
  

• 重启Jupyter，选择新注册的内核再次测试GPU检测。

验证步骤 #

在Jupyter Notebook中先执行以下代码，确认环境变量是否加载成功：

import os
print("PATH:", os.environ.get("PATH", ""))
print("LD_LIBRARY_PATH:", os.environ.get("LD_LIBRARY_PATH", ""))

若输出包含/home/user/cuda-11.2/bin和/home/user/cuda-11.2/lib64，说明路径配置生效，再重新检测GPU。

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