在 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

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

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

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

> wsl --set-default-version 2

新版：

直接执行 wsl --install

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

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

Step 3 检查安装成功

方法一：

> wsl -t -v

方法二:

> bash

安装 miniconda

进入 Ubuntu, 更新系统

$ sudo apt-get update$ sudo apt-get upgrade

下载 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

更新 conda

$ conda update conda

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

$ conda create -n python3.9 python=3.9

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

测试代码

import torch
if __name__ == '__main__':    print(torch.__version__)    print(torch.cuda.is_available())

结语

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