如何让校园Tambora服务器上的JupyterLab启用GPU运行？

一、先确认服务器GPU状态 #

• 在JupyterLab终端执行nvidia-smi，查看GPU是否存在、驱动版本及占用情况：
  • 无输出：说明服务器无GPU或驱动未安装，直接联系管理员
  • 有输出但无运行进程：GPU可用，需配置环境

二、针对Flax代码的GPU配置（核心，因为你用的是Flax框架） #

Flax基于JAX，需确保JAX启用GPU：

1. 检查JAX版本：
   在JupyterLab代码单元格执行：

   !pip list | grep jax
   

   若jaxlib后缀带cpu，说明是CPU版，需卸载后重装GPU版：

   !pip uninstall -y jax jaxlib
   !pip install "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
   

   （注：根据服务器CUDA版本调整，比如CUDA 12.x用jax[cuda12_pip]）

2. 验证GPU是否启用：

   import jax
   print(jax.devices())
   

   输出包含GpuDevice则说明GPU已生效

三、JupyterLab环境配置 #

• 切换到对应kernel：通过Kernel > Change Kernel选择安装了GPU版JAX/Flax的虚拟环境或conda环境
• 指定使用的GPU（多GPU场景）：

  import os
  os.environ['CUDA_VISIBLE_DEVICES'] = '0'  # 0为GPU编号，根据nvidia-smi结果调整
  

四、TensorFlow检测GPU问题的补充 #

如果你需要同时用TensorFlow：

1. 检查CUDA与TensorFlow版本匹配性，比如TensorFlow 2.10对应CUDA 11.2、cuDNN 8.1
2. 执行nvcc --version确认CUDA版本，安装对应TensorFlow版本
3. 设置环境变量：

   !export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH
   

五、最后方案 #

若以上操作均无效，联系校园服务器管理员，确认GPU是否对普通用户开放，或是否存在硬件/驱动故障

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