WSL环境下安装支持GPU的tensorflow踩坑实录

发布于 2025-01-01
更新于 2025-01-22

目标:希望在WSL环境下安装支持GPU的tensorflow。

注意版本间的匹配关系

TensorFlow 对 Python

TensorFlow 对 Python 版本和 CUDA/cuDNN 版本有严格的兼容性要求。
根据 TensorFlow 官方文档,以下是 TensorFlow 2.x 支持的 Python 版本:

TensorFlow 版本 支持的 Python 版本
TensorFlow 2.15 3.9, 3.10, 3.11
TensorFlow 2.14 3.9, 3.10, 3.11
TensorFlow 2.13 3.8, 3.9, 3.10, 3.11
TensorFlow 2.12 3.8, 3.9, 3.10, 3.11
TensorFlow 2.11 3.7, 3.8, 3.9, 3.10

推荐的 Python 版本

  • Python 3.9Python 3.10 是目前最稳定的选择,广泛支持 TensorFlow 2.x 的最新版本。
  • 如果需要使用较旧的 TensorFlow 版本(如 2.11),可以选择 Python 3.8
  • Python 3.12 目前尚未被 TensorFlow 官方支持,因此不推荐使用。

CUDA/cuDNN 版本匹配

TensorFlow 的 GPU 支持依赖于 CUDA 和 cuDNN。以下是 TensorFlow 2.x 的 CUDA/cuDNN 兼容性:

TensorFlow 版本 CUDA 版本 cuDNN 版本
TensorFlow 2.15 CUDA 12.0 cuDNN 8.9
TensorFlow 2.14 CUDA 11.8 cuDNN 8.6
TensorFlow 2.13 CUDA 11.8 cuDNN 8.6
TensorFlow 2.12 CUDA 11.8 cuDNN 8.6
TensorFlow 2.11 CUDA 11.2 cuDNN 8.1
  • Python 3.10 + TensorFlow 2.15 + CUDA 12.0 + cuDNN 8.9(最新稳定组合)
  • Python 3.9 + TensorFlow 2.14 + CUDA 11.8 + cuDNN 8.6(广泛测试的组合)

对已经安装的验证

验证cuda的安装

sh
nvidia-smi

查看显示驱动支持的 CUDA 版本

sh
nvcc --version
# 或者
nvcc -V

查看输出版本信息,输出的示例如下

txt
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2023 NVIDIA Corporation
Built on Mon_Apr__3_17:36:15_Pacific_Daylight_Time_2023
Cuda compilation tools, release 12.1, V12.1.105
Build cuda_12.1.r12.1/compiler.32688072_0

最终验证,在python环境中检查是否支持GPU

python
import tensorflow as tf
print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')))

输出示例如下

txt
Num GPUs Available:  1

安装步骤

第一步,去NVIDIA官网,按官网指示安装驱动。
在Game Ready Driver 和 Studio Drive两种当中,鉴于我们进行的是tensorflow开发,所以选择Studio Drive更合适。
在安装向导中选择 自定义安装(Custom Installation)
勾选 执行干净安装(Perform a Clean Installation) (这个选项会删除旧版本的设置和文件,但保留基本驱动组件)。。
按照安装程序提示完成安装。
打开命令提示符,输入以下命令,查看显卡驱动是否安装成功。

sh
nvidia-smi

第二步,安装CUDA Toolkit
https://developer.nvidia.com/cuda-toolkit

CUDA Toolkit安装CUDA Toolkit安装

sh
wget https://developer.download.nvidia.com/compute/cuda/repos/wsl-ubuntu/x86_64/cuda-wsl-ubuntu.pin
sudo mv cuda-wsl-ubuntu.pin /etc/apt/preferences.d/cuda-repository-pin-600
wget https://developer.download.nvidia.com/compute/cuda/12.6.3/local_installers/cuda-repo-wsl-ubuntu-12-6-local_12.6.3-1_amd64.deb
sudo dpkg -i cuda-repo-wsl-ubuntu-12-6-local_12.6.3-1_amd64.deb
sudo cp /var/cuda-repo-wsl-ubuntu-12-6-local/cuda-*-keyring.gpg /usr/share/keyrings/
sudo apt-get update
sudo apt-get -y install cuda-toolkit-12-6

按照上面官网的提示安装完成后还不够,还需要配置环境变量

sh
echo 'export PATH=/usr/local/cuda-12.6/bin:$PATH' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda-12.6/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc
source ~/.bashrc

配置完成后,删去安装包.deb文件

sh
rm cuda-repo-wsl-ubuntu-12-6-local_12.6.3-1_amd64.deb

检查cuda是否安装成功

sh
nvcc --version

查看输出信息,确认无误。

接下来创建虚拟环境,安装tensorflow即可

sh
python3 -m venv tf_env
source tf_env/bin/activate
pip install --upgrade pip
pip install tensorflow

这里tensorflow会自动支持gpu,不需要按照旧版本安装tensorflow-gpu。
验证是否支持gpu

python
import tensorflow as tf
print("Num GPUs Available: ", len(tf.config.list_physical_devices('GPU')))

如仍有其他问题,咨询deepseek或chatgpt即可,注意需要将提示词描述清晰、完整;或者卸载wsl重装。