企业官网建设流程全解析

从零搭建WSL2 AI炼丹炉：手把手配置PyTorch（CUDA 11.6）+ TensorRT 8.4推理环境

在Windows 11上通过WSL2搭建完整的AI开发环境，已经成为越来越多机器学习工程师的首选方案。这种组合既能享受Windows系统的易用性，又能获得接近原生Linux的性能表现。本文将带你从零开始，在WSL2的Ubuntu环境中配置一套完整的PyTorch训练和TensorRT推理环境，特别针对CUDA 11.6和TensorRT 8.4版本进行优化配置。

1. 环境准备与基础配置

1.1 Windows 11与WSL2基础设置

首先确保你的Windows 11系统已经启用WSL2功能。在管理员权限的PowerShell中运行以下命令：

wsl --install wsl --set-default-version 2

安装完成后，从Microsoft Store获取Ubuntu 22.04 LTS版本。建议选择LTS版本以获得长期稳定的支持。安装后，通过以下命令确认WSL版本：

wsl -l -v

如果显示为版本1，可以通过以下命令升级：

wsl --set-version Ubuntu-22.04 2

1.2 NVIDIA驱动检查与配置

WSL2的一个优势是可以直接使用Windows主机安装的NVIDIA驱动，无需在Ubuntu中重复安装。首先在Windows中确认驱动版本：

1. 打开NVIDIA控制面板
2. 点击"系统信息"
3. 查看"驱动程序版本"和"CUDA版本"

记录下显示的CUDA版本（例如11.6），这将是后续安装的基础。在WSL的Ubuntu终端中，运行：

nvidia-smi

如果正常显示GPU信息，说明驱动配置正确。如果遇到问题，可能需要更新Windows端的NVIDIA驱动。

2. CUDA 11.6环境搭建

2.1 CUDA Toolkit安装

CUDA是深度学习开发的基础，版本选择至关重要。根据前面nvidia-smi显示的版本，我们选择对应的CUDA 11.6进行安装：

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/11.6.2/local_installers/cuda-repo-wsl-ubuntu-11-6-local_11.6.2-1_amd64.deb sudo dpkg -i cuda-repo-wsl-ubuntu-11-6-local_11.6.2-1_amd64.deb sudo apt-key add /var/cuda-repo-wsl-ubuntu-11-6-local/7fa2af80.pub sudo apt-get update sudo apt-get -y install cuda

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

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

验证安装：

nvcc -V

2.2 cuDNN 8.4安装

cuDNN是NVIDIA提供的深度神经网络加速库，需要与CUDA版本严格匹配。下载cuDNN 8.4.0 for CUDA 11.6的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

3. TensorRT 8.4部署与配置

3.1 TensorRT安装

TensorRT是NVIDIA的高性能推理引擎，同样需要版本匹配。下载TensorRT 8.4.1 for CUDA 11.6的tar包后：

tar -xzvf TensorRT-8.4.1.5.Linux.x86_64-gnu.cuda-11.6.cudnn8.4.tar.gz

添加环境变量：

echo 'export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:~/TensorRT-8.4.1.5/lib' >> ~/.bashrc source ~/.bashrc

安装Python包：

cd TensorRT-8.4.1.5/python pip install tensorrt-8.4.1.5-cp39-none-linux_x86_64.whl

3.2 TensorRT验证

运行MNIST示例验证安装：

cd TensorRT-8.4.1.5/samples/sampleOnnxMNIST make cd ../../bin/ ./sample_onnx_mnist

Python环境验证：

import tensorrt as trt print(trt.__version__) assert trt.Builder(trt.Logger())

4. PyTorch环境配置

4.1 Conda环境创建

建议使用conda管理Python环境：

conda create -n pytorch python=3.9 conda activate pytorch

4.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/cu116

验证安装：

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

4.3 环境一致性检查

确保所有组件的CUDA版本一致：

5. 开发环境优化

5.1 VS Code远程开发配置

1. 在Windows上安装VS Code和"Remote - WSL"扩展
2. 在WSL终端中运行code .启动VS Code
3. 安装Python、Pylance、Jupyter等扩展

5.2 Jupyter Notebook配置

在conda环境中安装Jupyter：

pip install jupyter

配置密码并生成配置文件：

jupyter notebook --generate-config jupyter notebook password

启动Notebook：

jupyter notebook --ip=0.0.0.0 --no-browser

5.3 性能优化技巧

1. 将数据集放在WSL的文件系统中（/home/username）而非Windows挂载目录（/mnt/c）
2. 增加WSL内存限制，在%USERPROFILE%\.wslconfig中添加：

[wsl2] memory=16GB swap=8GB localhostForwarding=true

3. 使用torch.backends.cudnn.benchmark = True启用cuDNN自动调优

6. 完整工作流验证

6.1 训练示例

创建一个简单的PyTorch训练脚本train.py：

import torch import torch.nn as nn import torch.optim as optim # 定义简单模型 model = nn.Sequential( nn.Linear(10, 50), nn.ReLU(), nn.Linear(50, 2) ).cuda() # 模拟数据 inputs = torch.randn(100, 10).cuda() targets = torch.randint(0, 2, (100,)).cuda() # 训练循环 criterion = nn.CrossEntropyLoss() optimizer = optim.Adam(model.parameters()) for epoch in range(10): optimizer.zero_grad() outputs = model(inputs) loss = criterion(outputs, targets) loss.backward() optimizer.step() print(f'Epoch {epoch}, Loss: {loss.item()}')

运行脚本验证训练功能：

python train.py

6.2 模型导出与TensorRT推理

将PyTorch模型导出为ONNX格式：

import torch dummy_input = torch.randn(1, 10).cuda() torch.onnx.export(model, dummy_input, "model.onnx", input_names=["input"], output_names=["output"])

使用TensorRT进行推理：

import tensorrt as trt logger = trt.Logger(trt.Logger.WARNING) builder = trt.Builder(logger) network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) parser = trt.OnnxParser(network, logger) with open("model.onnx", "rb") as f: parser.parse(f.read()) config = builder.create_builder_config() config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30) serialized_engine = builder.build_serialized_network(network, config) with open("model.engine", "wb") as f: f.write(serialized_engine)

7. 常见问题解决

7.1 CUDA版本不匹配

如果遇到类似错误：

CUDA error: no kernel image is available for execution on the device

解决方案：

1. 确认所有组件的CUDA版本一致
2. 重新安装匹配版本的PyTorch
3. 检查torch.cuda.is_available()返回是否为True

7.2 TensorRT初始化失败

如果assert trt.Builder(trt.Logger())失败：

1. 检查LD_LIBRARY_PATH是否包含TensorRT的lib路径
2. 确认Python环境中tensorrt包的版本与安装的TensorRT一致
3. 尝试重新安装TensorRT的Python包

7.3 WSL2内存不足

如果训练过程中内存不足：

1. 增加.wslconfig中的内存限制
2. 减小batch size
3. 使用梯度累积等技术

8. 进阶配置

8.1 多GPU训练支持

配置多GPU训练环境：

import torch.distributed as dist from torch.nn.parallel import DistributedDataParallel as DDP dist.init_process_group("nccl") model = DDP(model.cuda())

8.2 混合精度训练

启用自动混合精度训练：

from torch.cuda.amp import GradScaler, autocast scaler = GradScaler() with autocast(): outputs = model(inputs) loss = criterion(outputs, targets) scaler.scale(loss).backward() scaler.step(optimizer) scaler.update()

8.3 自定义CUDA扩展

编写并编译自定义CUDA操作：

// my_kernel.cu #include <torch/extension.h> torch::Tensor my_add(torch::Tensor a, torch::Tensor b) { return a + b; } PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("my_add", &my_add, "Add two tensors"); }

编译并安装：

python -m torch.utils.cpp_extension compile --cuda-archs=80 --verbose my_kernel.cu