wsl安装pytorch

文章目录

wsl安装pytorch
- [1. 安装之前](#1. 安装之前)
- [2. 安装CUDA Toolkit](#2. 安装CUDA Toolkit)
- - [2.1 确保 WSL2 已启用并更新：](#2.1 确保 WSL2 已启用并更新：)
  - [2.2 安装WSL-Ubuntu 专用包](#2.2 安装WSL-Ubuntu 专用包)
  - - [步骤 1：清理旧 GPG key（防止报错）](#步骤 1：清理旧 GPG key（防止报错）)
    - [步骤 2：添加 WSL-Ubuntu 专用仓库并安装最新 CUDA Toolkit](#步骤 2：添加 WSL-Ubuntu 专用仓库并安装最新 CUDA Toolkit)
  - [2.3 安装完成后配置环境变量](#2.3 安装完成后配置环境变量)
  - [2.4 验证安装是否成功](#2.4 验证安装是否成功)
[3. 安装pytorch](#3. 安装pytorch)
[额外：wsl安装NVIDIA Container Toolkit](#额外：wsl安装NVIDIA Container Toolkit)

1. 安装之前

wsl中安装cuda之前，需要先在Windows中安装NVIDIA驱动 (如Game Ready)。安装之后，在WSL2内部绝对不能安装Linux的NVIDIA驱动，以免而这发生冲突。

这一步相对简单，直接搜英伟达驱动安装就可以了，这里不做赘述。

2. 安装CUDA Toolkit

这是为了给wsl提供nvcc编译器、CUDA运行时库和开发工具。

2.1 确保 WSL2 已启用并更新：

bash 复制代码

wsl --update
wsl --version

# 重启wsl让更新生效
wsl --shutdown

重新打开wsl ubuntu，先更新系统：

bash 复制代码

sudo apt update && sudo apt upgrade -y

2.2 安装WSL-Ubuntu 专用包

专用包是官方推荐的，避免驱动冲突。

步骤 1：清理旧 GPG key（防止报错）

bash 复制代码

# 步骤 1：清理旧 GPG key（防止报错）
sudo apt-key del 7fa2af80 2>/dev/null || true

步骤 2：添加 WSL-Ubuntu 专用仓库并安装最新 CUDA Toolkit

专用仓库是为了避免安装 Linux 驱动。

目前最新稳定版通常是 CUDA 12.x 或 13.x（建议根据你的驱动支持选择）。以下以 CUDA 13.1 为例，请替换为你想要的版本）：

bash 复制代码

# 下载 pin 文件
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 13.1 local repo（如果有具体版本号请替换）
wget https://developer.download.nvidia.com/compute/cuda/13.1.0/local_installers/cuda-repo-wsl-ubuntu-13-1-local_13.1.0-1_amd64.deb

# 安装 repo
sudo dpkg -i cuda-repo-wsl-ubuntu-13-1-local_13.1.0-1_amd64.deb

# 添加 keyring
sudo cp /var/cuda-repo-wsl-ubuntu-13-1-local/cuda-*-keyring.gpg /usr/share/keyrings/

# 更新并安装 Toolkit（只安装 toolkit，不要带 drivers）
sudo apt update
sudo apt install -y cuda-toolkit-13-1  # 其他版本自行替换，如CUDA 12.6：cuda-toolkit-12-6

2.3 安装完成后配置环境变量

bash 复制代码

# 添加到 ~/.bashrc（永久生效）
echo 'export PATH=/usr/local/cuda/bin:$PATH' >> ~/.bashrc
echo 'export LD_LIBRARY_PATH=/usr/local/cuda/lib64:$LD_LIBRARY_PATH' >> ~/.bashrc

source ~/.bashrc

2.4 验证安装是否成功

bash 复制代码

# 1. 检查 nvcc（CUDA 编译器）
nvcc --version

# 2. 检查 GPU 信息（应该显示 Windows 驱动支持的 CUDA Version）
nvidia-smi

可选额外步骤：编译并运行官方 sample

bash 复制代码

# 先确保必要的构建工具已安装
sudo apt install -y build-essential cmake git

# 下载最新 CUDA Samples
cd ~
git clone https://github.com/NVIDIA/cuda-samples.git

# 编译 deviceQuery
cd ~/cuda-samples/Samples/1_Utilities/deviceQuery
mkdir -p build && cd build
cmake ..
make

# 运行
./deviceQuery

如果 deviceQuery 显示 Result = PASS，则安装成功

3. 安装pytorch

首先查看自己的gpu信息

bash 复制代码

nvidia-smi

然后问一下ai应该用什么版本的cuda，创建虚拟环境后pip安装。这里可能会非常慢，所以使用清华pypi，能提速80%。

bash 复制代码

python3 -m venv pytorch_env
source pytorch_env/bin/activate
pip install --upgrade pip

pip install torch torchvision torchaudio -i https://pypi.tuna.tsinghua.edu.cn/simple

注：也可以将镜像源添加到配置文件中，以后自动使用。创建文件~/.pip/pip.conf，添加下面两行即可：

conf 复制代码

[global]
index-url = https://pypi.tuna.tsinghua.edu.cn/simple

最后验证

bash 复制代码

python3

import torch
print(torch.__version__)
print(torch.cuda.is_available())          # 应该返回 True
print(torch.cuda.get_device_name(0))      # 应该显示 RTX 5070 Ti
print(torch.cuda.get_device_capability(0)) # 应该看到 (12, 0)

WSL+RTX50系列显卡，很多用户成功安装的版本是cu128，也是上述范例中写的。但如果是blackwell架构，安装后出现"not compatible"警告，可以尝试nightly版本：

bash 复制代码

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

额外：wsl安装NVIDIA Container Toolkit

NVIDIA Container Toolkit 是一个工具集，使用户能够在容器化环境（如 Docker 或 Kubernetes）中高效使用 NVIDIA GPU 进行加速计算。它提供了一整套工具，使用户可以轻松将 NVIDIA GPU 的计算能力集成到容器中，用于深度学习、科学计算等领域的工作负载。

NVIDIA Container Toolkit 允许在容器中使用 GPU 资源，从而使得基于容器的应用可以在各种平台上轻松运行并利用 GPU 进行加速。
容器运行时（如 Docker）通过 NVIDIA Container Runtime 将 GPU 计算资源暴露给容器中的应用。
通过与 Docker 的无缝集成，用户可以通过 docker run 命令启动使用 GPU 的容器，只需要添加 --gpus 参数即可，如：docker run --gpus all。
工具集还支持将多个 GPU 映射到容器中，并控制每个容器可以使用的 GPU 数量或具体的 GPU 设备。

首先，参照官方安装指导安装。

然后配置docker

bash 复制代码

sudo nvidia-ctk runtime configure --runtime=docker

在windows中重启docker（右键任务栏中的图标，restart）。

测试是否安装成功：

bash 复制代码

sudo docker run --rm --gpus all ubuntu nvidia-smi