从零搭建AI开发环境:在Win11的WSL Ubuntu里配置PyTorch(CUDA 11.6)完整流程
从零搭建AI开发环境:Win11 WSL Ubuntu中PyTorch(CUDA 11.6)全流程配置指南
在Windows 11上通过WSL 2运行Ubuntu进行AI开发,已经成为越来越多开发者的选择。这种组合既能享受Windows的易用性,又能获得Linux环境的高效开发体验。本文将带你从零开始,在Win11的WSL Ubuntu中配置完整的PyTorch开发环境,特别针对CUDA 11.6版本进行详细说明。
1. 环境准备与基础配置
1.1 WSL 2与Ubuntu安装
首先确保你的Windows 11系统已启用WSL 2功能。在管理员权限的PowerShell中执行以下命令:
wsl --install wsl --set-default-version 2安装完成后,从Microsoft Store获取Ubuntu 20.04或22.04 LTS版本。首次启动Ubuntu时会自动完成初始化设置,记得创建非root用户账户。
1.2 NVIDIA驱动检查
WSL 2的一个优势是可以直接使用Windows主机安装的NVIDIA驱动。在Windows中确认已安装最新驱动后,在Ubuntu终端中运行:
nvidia-smi正常输出应显示类似以下信息:
+-----------------------------------------------------------------------------+ | NVIDIA-SMI 525.85.12 Driver Version: 525.85.12 CUDA Version: 12.0 | |-------------------------------+----------------------+----------------------+注意:WSL中的CUDA版本显示可能高于实际安装版本,这是正常现象
2. CUDA工具包安装
2.1 选择正确的CUDA版本
针对PyTorch稳定支持的情况,我们选择CUDA 11.6版本。首先添加NVIDIA官方仓库:
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-6002.2 安装CUDA 11.6
执行以下命令完成安装:
sudo apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/wsl-ubuntu/x86_64/3bf863cc.pub sudo add-apt-repository "deb https://developer.download.nvidia.com/compute/cuda/repos/wsl-ubuntu/x86_64/ /" sudo apt-get update sudo apt-get -y install cuda-11-6安装完成后,验证CUDA编译器版本:
nvcc --version2.3 环境变量配置
将以下内容添加到~/.bashrc文件末尾:
export PATH=/usr/local/cuda-11.6/bin${PATH:+:${PATH}} export LD_LIBRARY_PATH=/usr/local/cuda-11.6/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}} export CUDA_HOME=/usr/local/cuda-11.6执行source ~/.bashrc使配置生效。
3. cuDNN与TensorRT安装
3.1 cuDNN 8.4安装
从NVIDIA开发者网站下载对应CUDA 11.6的cuDNN 8.4 tar包后,执行:
tar -xvf cudnn-linux-x86_64-8.4.0.27_cuda11.6-archive.tar.xz sudo cp cudnn-*-archive/include/cudnn*.h /usr/local/cuda/include sudo cp -P cudnn-*-archive/lib/libcudnn* /usr/local/cuda/lib64 sudo chmod a+r /usr/local/cuda/include/cudnn*.h /usr/local/cuda/lib64/libcudnn*验证安装:
cat /usr/local/cuda/include/cudnn_version.h | grep CUDNN_MAJOR -A 23.2 TensorRT 8.4安装
下载对应CUDA 11.6的TensorRT 8.4 GA tar包后:
tar -xzvf TensorRT-8.4.1.5.Linux.x86_64-gnu.cuda-11.6.cudnn8.4.tar.gz添加环境变量到~/.bashrc:
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/path/to/TensorRT-8.4.1.5/lib安装Python包:
cd TensorRT-8.4.1.5/python pip install tensorrt-8.4.1.5-cp39-none-linux_x86_64.whl4. PyTorch环境配置
4.1 创建Python虚拟环境
推荐使用conda或venv创建独立环境:
conda create -n pytorch python=3.9 conda activate pytorch4.2 安装PyTorch及相关库
针对CUDA 11.6安装对应版本的PyTorch:
pip install torch==1.12.0+cu116 torchvision==0.13.0+cu116 torchaudio==0.12.0 --extra-index-url https://download.pytorch.org/whl/cu1164.3 环境验证
创建测试脚本test_gpu.py:
import torch print(f"PyTorch版本: {torch.__version__}") print(f"CUDA可用: {torch.cuda.is_available()}") print(f"GPU数量: {torch.cuda.device_count()}") print(f"当前GPU: {torch.cuda.current_device()}") print(f"GPU名称: {torch.cuda.get_device_name(0)}")运行结果应显示CUDA可用且能正确识别GPU设备。
5. 常见问题排查
5.1 CUDA版本不匹配
如果遇到CUDA相关错误,首先检查各组件版本一致性:
| 组件 | 推荐版本 | 验证命令 |
|---|---|---|
| CUDA | 11.6 | nvcc --version |
| cuDNN | 8.4 | cat /usr/local/cuda/include/cudnn_version.h |
| TensorRT | 8.4 GA | python -c "import tensorrt; print(tensorrt.__version__)" |
| PyTorch | 1.12.0+cu116 | python -c "import torch; print(torch.__version__)" |
5.2 共享内存问题
在WSL 2中运行大型模型可能遇到共享内存不足的问题,解决方法:
sudo mount -t tmpfs -o size=8G tmpfs /dev/shm可将此命令添加到~/.bashrc中自动执行。
5.3 GPU内存管理
WSL 2中的GPU内存分配策略与原生Linux不同,可通过以下方式优化:
# 在PyTorch代码开始处添加 torch.cuda.set_per_process_memory_fraction(0.8) # 限制进程使用80%的GPU内存6. 开发环境优化
6.1 VS Code远程开发配置
- 安装VS Code的"Remote - WSL"扩展
- 在WSL终端中执行
code .启动VS Code - 安装Python、Pylance等扩展
6.2 Jupyter Notebook配置
在虚拟环境中安装并配置Jupyter:
pip install jupyterlab jupyter lab --ip=0.0.0.0 --port=8888 --no-browser在Windows浏览器中访问localhost:8888即可使用。
6.3 性能调优建议
- 将项目文件存储在WSL文件系统中(
\\wsl$\Ubuntu-20.04\home\user) - 禁用Windows Defender对WSL目录的实时保护
- 在
%UserProfile%\.wslconfig中添加:
[wsl2] memory=16GB processors=8 localhostForwarding=true经过以上步骤,你已经拥有了一个完整的AI开发环境。在实际项目中,这种配置可以流畅运行大多数计算机视觉和自然语言处理模型。
