从零搭建WSL2深度学习环境：用Anaconda管理CUDA 11.6、PyTorch和TensorRT

在Windows 11上通过WSL2搭建深度学习开发环境，已经成为越来越多开发者的选择。这种方案既能享受Windows系统的易用性，又能获得Linux环境下高效的开发体验。对于需要同时处理多个项目的机器学习工程师来说，如何在同一台机器上管理不同版本的CUDA、PyTorch和TensorRT，同时保持环境隔离和项目独立性，是一个常见且棘手的问题。

本文将详细介绍如何在WSL2的Ubuntu系统中，使用Anaconda创建隔离的Python环境，并正确配置CUDA 11.6、PyTorch和TensorRT的开发环境。不同于简单的安装教程，我们更关注工程实践中的环境隔离和版本管理问题，帮助你在同一台机器上为不同项目维护独立的深度学习环境。

1. 环境准备与基础配置

在开始之前，确保你的系统满足以下基本要求：

• Windows 11 21H2或更高版本
• 已启用WSL2功能并安装Ubuntu发行版（建议20.04 LTS或22.04 LTS）
• NVIDIA显卡驱动已正确安装在Windows主机上（无需在WSL中重复安装）
• 至少50GB的可用磁盘空间（深度学习环境占用空间较大）

首先，我们需要验证WSL2和NVIDIA驱动的准备工作：

# 在Windows PowerShell中检查WSL版本
wsl --list --verbose

# 在WSL Ubuntu终端中检查NVIDIA驱动
nvidia-smi

如果nvidia-smi命令能够正确显示GPU信息，说明NVIDIA驱动已经正确配置。接下来，我们需要在WSL中安装一些基础工具：

sudo apt update && sudo apt upgrade -y
sudo apt install -y build-essential cmake git wget unzip

2. Miniconda安装与环境管理

Anaconda/Miniconda是管理Python环境的理想工具，特别适合需要同时维护多个项目的开发者。我们选择安装更轻量级的Miniconda：

# 下载并安装Miniconda
wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
bash Miniconda3-latest-Linux-x86_64.sh -b -p $HOME/miniconda

# 初始化conda
~/miniconda/bin/conda init bash
source ~/.bashrc

安装完成后，我们可以创建专为深度学习项目设计的conda环境。假设我们有两个项目：一个使用PyTorch 1.12，另一个使用TensorFlow：

# 创建PyTorch项目环境
conda create -n pytorch_project python=3.9 -y

# 创建TensorFlow项目环境
conda create -n tensorflow_project python=3.8 -y

3. 系统级CUDA工具链安装

虽然conda可以管理CUDA工具链，但对于需要高性能计算的场景，建议在系统级别安装CUDA，然后在conda环境中链接这些库。以下是CUDA 11.6的安装步骤：

# 添加NVIDIA CUDA仓库
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

# 安装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

安装完成后，配置环境变量：

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

验证CUDA安装：

nvcc --version

4. cuDNN和TensorRT的安装与配置

cuDNN和TensorRT是深度学习加速的重要组件。由于它们与CUDA版本紧密相关，我们需要确保版本匹配。

4.1 cuDNN安装

下载对应CUDA 11.6的cuDNN 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 2

4.2 TensorRT安装

TensorRT的安装过程类似：

tar -xzvf TensorRT-8.4.1.5.Linux.x86_64-gnu.cuda-11.6.cudnn8.4.tar.gz
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$(pwd)/TensorRT-8.4.1.5/lib
echo 'export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:'$(pwd)'/TensorRT-8.4.1.5/lib' >> ~/.bashrc

# 安装Python包
cd TensorRT-8.4.1.5/python
pip install tensorrt-8.4.1.5-cp39-none-linux_x86_64.whl

5. 在conda环境中安装深度学习框架

现在，我们可以在隔离的conda环境中安装PyTorch和TensorFlow，同时确保它们使用系统安装的CUDA工具链。

5.1 PyTorch环境配置

激活PyTorch项目环境并安装框架：

conda activate pytorch_project
pip install torch==1.12.0+cu116 torchvision==0.13.0+cu116 torchaudio==0.12.0 --extra-index-url https://download.pytorch.org/whl/cu116

验证PyTorch是否能正确使用GPU：

import torch
print(torch.__version__)
print(torch.cuda.is_available())
print(torch.cuda.get_device_name(0))

5.2 TensorFlow环境配置

对于TensorFlow项目环境：

conda activate tensorflow_project
pip install tensorflow-gpu==2.9.1

验证TensorFlow GPU支持：

import tensorflow as tf
print(tf.__version__)
print(tf.config.list_physical_devices('GPU'))

6. 环境变量管理与项目隔离

为了确保不同项目使用正确的库路径，我们需要为每个conda环境设置特定的环境变量。创建环境激活时的钩子脚本：

# 为PyTorch环境创建激活脚本
mkdir -p ~/miniconda/envs/pytorch_project/etc/conda/activate.d
echo 'export LD_LIBRARY_PATH=/usr/local/cuda-11.6/lib64:$LD_LIBRARY_PATH' > ~/miniconda/envs/pytorch_project/etc/conda/activate.d/env_vars.sh

# 为TensorFlow环境创建激活脚本
mkdir -p ~/miniconda/envs/tensorflow_project/etc/conda/activate.d
echo 'export LD_LIBRARY_PATH=/usr/local/cuda-11.6/lib64:$LD_LIBRARY_PATH' > ~/miniconda/envs/tensorflow_project/etc/conda/activate.d/env_vars.sh

这样，每次激活conda环境时，都会自动设置正确的库路径，确保框架能够找到系统安装的CUDA库。

7. 常见问题排查与性能优化

在配置过程中可能会遇到各种问题，以下是一些常见问题的解决方法：

问题1：PyTorch找不到CUDA

解决方案：确保conda环境中没有安装cudatoolkit包，这可能会与系统CUDA冲突。使用conda list检查并移除冲突的包。

问题2：TensorRT无法初始化

解决方案：检查LD_LIBRARY_PATH是否包含TensorRT库路径，并确保TensorRT版本与CUDA版本匹配。

性能优化建议：

1. 在WSL2配置文件中增加内存和CPU核心分配：

   # 在Windows的%USERPROFILE%/.wslconfig中添加
   [wsl2]
   memory=16GB
   processors=8
   

2. 使用WSL2的GPU计算特性需要Windows 11 22H2或更高版本，确保系统已更新。

3. 对于大型数据集，考虑将数据存储在WSL2文件系统之外，以提高I/O性能。

通过以上步骤，你已经成功在WSL2中搭建了一个灵活、高效的深度学习开发环境，能够同时支持多个项目的独立开发需求。这种配置方式既保持了环境的整洁，又能充分利用系统资源，是机器学习工程师的理想工作环境配置。