在 WSL2 上部署 PyTorch

开发项目需要使用AI，但是我习惯用 VMware 来开发各类项目，于是查到《Enable NVIDIA CUDA on WSL》可以实现 CUDA + VGPU 的功能。但是国内对这类资料较少，我就想顺便补个空白——在 WSL2 上部署 PyTorch。

前提

AI开发的工作者无非两种：

1. 在 Ubuntu 或 Debian 上直接用Nouveau部署带 CUDA 的AI框架；
2. 直接在 Windows 上用 Nvidia 官方的驱动，部署带CUDA的AI框架；

第一种方法，2019年我尝试过安装 Nvidia 的驱动，除了 Ubuntu 支持的最佳，其他发行版也就一般，希望这些社区能努力一下（我当时为14年的笔记本尝试各种发行版的安装测试），即使今年 Nvidia 突然声称为 Linux 搞官方驱动。第二种方法，虽然安装驱动方便多了，但是 Windows 一般用 MSVC，而多数AI框架都是基于 GCC 开发的。当然也可以用 MinGW 或 cygwin 来实现 Windows 下的 GCC。

作为一个有洁癖的开发工程师，我一般都是在 VM 中开发的，优缺点懂得都懂。有些习惯就像“On your left”改不了了。VMware 只有 ESXI 支持 V-Gpu; VirtualBox 据我所知没有显卡直通，有也是 Intel 集显。这里非常感谢微软，在致力于拥抱 Docker 的时候，还开发了 WSL2。毕竟之前就发布过 Hyper-V，但是我觉得 WSL 相比于 Hyper-V 更加轻量，且对于 Windows 的资源整合度更高，可以在PS上直接调用 WSL 的 Bash。

安装 WSL2

Step 1 更新Windows 10 21H2

Step 2 启用 WSL2

Windows 10

• x64系统需要Version 1903，Build 18362
• ARM64需要Version 2004，Build 19041
  或者Windows 11

旧版：

升级后还需要进行一些配置才可以使用 WSL2，首先要启用 Windows 子系统功能，使用管理员权限打开一个 PowerShell 窗口，输入以下命令，并重启系统：

> dism.exe /online /enable-feature /featurename:Microsoft-Windows-Subsystem-Linux /all /norestart
1

重启后，Windows 默认启用的是 WSL1，还需要再启用虚拟机平台功能，在 PowerShell 中输入以下命令，并再次重启系统：

> dism.exe /online /enable-feature /featurename:VirtualMachinePlatform /all /norestart
1

重启后，在 PowerShell 中输入以下命令，将 WSL 默认版本改为 WSL2:

> wsl --set-default-version 2
1

新版：

直接执行 wsl --install

确认控制面板的启动或关闭 Windows 功能

Linux 的发行版可以在 Microsoft Store中下载。推荐Ubuntu

Step 3 检查安装成功

方法一：

> wsl -t -v
1

方法二:

> bash
1

安装miniconda

进入Ubuntu, 更新系统

$ sudo apt-get update
$ sudo apt-get upgrade
1
2

下载miniconda

$ wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh
$ chmod +x Miniconda3-latest-Linux-x86_64.sh
$ ./Miniconda3-latest-Linux-x86_64.sh
1
2
3

更新conda

$ conda update conda
1

创建python的虚拟环境, 目前PyTorch支持到3.9

$ conda create -n python3.9 python=3.9
1

安装Nvidia驱动

从官网上下载，WSL上可以不必安装其他Nvidia驱动。

安装PyTorch

按照PyTorch官网提供的pip命令下载太慢。直接在 https://download.pytorch.org/whl/torch_stable.html下载

注意：CU116是CUDA 11.6；CP3.9是python3.9

pip安装

$ pip3 install ./torch-1.12.1+cu116-cp39-cp39-linux_x86_64.whl
$ pip3 install torch===1.12.1 torchvision torchaudio -f https://download.pytorch.org/whl/torch_stable.html
1
2

测试代码

import torch

if __name__ == '__main__':
    print(torch.__version__)
    print(torch.cuda.is_available())
1
2
3
4
5

结语

这种方式主要适合AI开发环境，而生产环境还是应该用Linux环境。一般云计算服务厂商都帮我们直接部署好了。当然有条件的也可以用Mac OSX + AMD显卡。