WSL Ubuntu 下使用 CUDA 进行 AI 推理和训练指南

WSL Ubuntu 下使用 CUDA 进行 AI 推理和训练指南

在 WSL Ubuntu 环境下配置和使用 CUDA 进行 AI 推理和训练的详细指南。

1. 环境准备

1.1 确认 WSL 版本

首先确保您使用的是 WSL2，因为只有 WSL2 支持 CUDA：

wsl --list --verbose

如果显示版本为 1，需要升级到 WSL2：

wsl --set-version Ubuntu 2

1.2 安装必要的驱动

1. 在 Windows 主机上安装最新的 NVIDIA 显卡驱动（支持 WSL）
2. 确保在 Windows 设置中启用了 WSL 2 GPU 支持

2. 在 WSL Ubuntu 中安装 CUDA 工具包

2.1 添加 NVIDIA 软件源

wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-keyring_1.0-1_all.deb
sudo dpkg -i cuda-keyring_1.0-1_all.deb
sudo apt-get update

2.2 安装 CUDA 工具包

sudo apt-get install -y cuda-toolkit-12-x

2.3 设置环境变量

将以下内容添加到 ~/.bashrc 或 ~/.zshrc（取决于您使用的 shell）：

export PATH=/usr/local/cuda/bin:$PATH
export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH

然后应用更改：

source ~/.bashrc  # 或 source ~/.zshrc

2.4 验证 CUDA 安装

nvcc --version
nvidia-smi

3. 安装 AI 框架

3.1 创建 Python 虚拟环境

sudo apt-get install -y python3-pip python3-venv
python3 -m venv ~/ai-env
source ~/ai-env/bin/activate

3.2 安装 PyTorch（支持 CUDA）

pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

3.3 安装 TensorFlow（支持 CUDA）

pip install tensorflow

3.4 验证框架 CUDA 支持

PyTorch 验证：

import torch
print("CUDA 是否可用:", torch.cuda.is_available())
print("CUDA 设备数量:", torch.cuda.device_count())
print("当前 CUDA 设备:", torch.cuda.current_device())
print("CUDA 设备名称:", torch.cuda.get_device_name(0))

TensorFlow 验证：

import tensorflow as tf
print("TensorFlow 版本:", tf.__version__)
print("GPU 是否可用:", tf.config.list_physical_devices('GPU'))

4. AI 推理示例

4.1 PyTorch 推理示例

创建文件 inference_pytorch.py：

import torch
from torchvision import models
import time

# 加载预训练模型
model = models.resnet50(pretrained=True)

# 将模型移至 GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
model.eval()

# 创建随机输入数据
dummy_in