WSL2使用Nvidia-Docker实现CUDA版本自由切换

众所周知，深度学习的环境往往非常麻烦，经常不同的项目所依赖的 torch、tensorflow 包对 CUDA 的版本也有不同的要求，Linux 下进行 CUDA 的管理比较麻烦，是一个比较头疼的问题。

随着 WSL2 对物理机显卡的支持，Nvidia-Docker 也提供了对容器显卡的支持。我们可以通过拉取不同的 Docker 镜像的方式来实现对容器内 CUDA、CUDNN 的自由切换，操作非常简易。

1. Win11 显卡驱动的安装

注意：WSL2中是不需要且不能安装任何显卡驱动的，它的显卡驱动完全依赖于 Win11 中的显卡驱动，因此我们只需要安装你显卡对应的 Win11 版本显卡驱动版本（必须是 Win11 版本的驱动），这个已经有很多教程了，这里就不赘述。如果你安装成功，可以在 Win11 的 cmd 中输入 nvidia-smi可以看到下图。

因为 WSL2 中的显卡驱动完全依赖于 Win11 的显卡驱动，因此在 WSL2 中输入 nvidia-smi 也可以看到相同驱动版本的输出。
请注意：这里的 nvidia-smi 能作用的范围，只作用于你 Win11 安装显卡驱动时所登录的那个用户名对应到 WSL2 中的用户名。比如我是在 Win11 (guosongyuan) 用户上安装的显卡驱动，那么我只能在 WSL2 的 gsy 用户状态下才能执行该 nvidia-smi 指令，root 用户执行该命令是不能生效的。

2. 安装 Docker 和 Nvidia-Docker

1. 安装 Docker 引擎可以参考文档：Docker 引擎官方安装教程；
2. 安装 Docker 引擎之后，就可以在其基础上安装 Nvidia-Docker 组件：Nvidia-Docker 安装教程。
   这两个步骤非常简单，如果看不懂英语的话直接用谷歌翻译就好。

3. 选择合适的 CUDA 和 CUDNN 的镜像

使用 Nvidia-Docker 的好处就在于，你不需要真的在 WSL 中安装 CUDA 和 CUDNN，这样就可以避免在配置不同项目环境时遇到的很麻烦的环境切换问题。我们只要每次遇到一个新的项目，拉取对应的 CUDA 和 CUDNN 版本即可，即插即用，不想用了直接删除对应的镜像和容器即可，跟删除软件一样方便。

这里以安装 CUDA 11.2.0 版本为例，我们来到 Docker 镜像市场：Docker HUB，在其中搜索关键字 nvidia/cuda，如下图。

点进入，在 Tags 中搜索对应的 CUDA 版本，注意同一个版本下对应三种不同的类型（devel、runtime、base），我们推荐安装 devel 版本，因为它的环境更齐全，我们这里因为 WSL2 是 Ubuntu 20.04 版本的，所以我们选择镜像的时候选择 ubuntu20.04 后缀的。
这里以 nvidia/cuda:11.2.0-cudnn8-devel-ubuntu20.04 镜像为例，通过 sudo docker pull nvidia/cuda:11.2.0-cudnn8-devel-ubuntu20.04 将镜像拉取下来。

拉取镜像之后，我们可以查看当前镜像中的显卡驱动、CUDA版本和 CUDNN 的版本。

1. 查看显卡驱动版本：sudo docker run --rm --gpus all nvidia/cuda:11.2.0-cudnn8-devel-ubuntu20.04 nvidia-smi
2. 查看 CUDA 版本：sudo docker run --rm --gpus all nvidia/cuda:11.2.0-cudnn8-devel-ubuntu20.04 nvcc -V
3. 查看 CUDNN 版本，因为镜像官方将 CUDA 和 CUDNN 进行了解耦合，因此我们需要分两步进行查询操作。首先通过 sudo docker run --rm --gpus all nvidia/cuda:11.2.0-cudnn8-devel-ubuntu20.04 whereis cudnn，看到 cudnn.h 所在路径 cudnn: /usr/include/cudnn.h。我们根据这个输出结果，把 cudnn.h 之前的 include 路径记住，查询该 include 下的 cudnn_verseion.h 文件：sudo docker run --rm --gpus all nvidia/cuda:11.2.0-cudnn8-devel-ubuntu20.04 cat /usr/include/cudnn_version.h | grep CUDNN_MAJOR -A 2，这样就能看到 CUDNN 的版本号了。
   

4. 利用拉取的镜像构建自己的镜像

我们拉取的镜像中只有最基础的 CUDA 和 CUDNN，还没有配置 Anaconda、换源、git 、pip 等常用工具，因此我们将这些可能用到的常用工具将其打包好。

为了构建镜像，我们在用户目录下创建一个名为 mkimage 的目录，在其中放入我们需要的三个内容： Anaconda3-5.2.0-Linux-x86_64.sh、Dockerfile、sources.list，其中 sources.list 是用来给 Ubuntu apt 换源用的。

sources.list 内容如下：

######################################
###### CONTENT for sources.list ######
######################################

deb http://mirrors.aliyun.com/ubuntu/ focal main restricted universe multiverse
deb-src http://mirrors.aliyun.com/ubuntu/ focal main restricted universe multiverse

deb http://mirrors.aliyun.com/ubuntu/ focal-security main restricted universe multiverse
deb-src http://mirrors.aliyun.com/ubuntu/ focal-security main restricted universe multiverse

deb http://mirrors.aliyun.com/ubuntu/ focal-updates main restricted universe multiverse
deb-src http://mirrors.aliyun.com/ubuntu/ focal-updates main restricted universe multiverse

deb http://mirrors.aliyun.com/ubuntu/ focal-proposed main restricted universe multiverse
deb-src http://mirrors.aliyun.com/ubuntu/ focal-proposed main restricted universe multiverse

deb http://mirrors.aliyun.com/ubuntu/ focal-backports main restricted universe multiverse
deb-src http://mirrors.aliyun.com/ubuntu/ focal-backports main restricted universe multiverse

Dockerfile内容如下：

####################################
###### CONTENT for Dockerfile ######
####################################

# Extends from father image [CHANGE WITH YOUR NEED]
FROM nvidia/cuda:11.2.0-cudnn8-devel-ubuntu20.04

# Set locale
ENV DEBIAN_FRONTEND noninteractive

# Change anaconda source
# ADD means copy file from host machine to containers
ADD sources.list /etc/apt/
ENV PATH /opt/conda/bin:$PATH

# Install basic dependencies
RUN rm /etc/apt/sources.list.d/cuda.list && \
	rm /etc/apt/sources.list.d/nvidia-ml.list
	
RUN apt-get update && apt-get install -y --no-install-recommends \
	bzip2 \
	g++ \
	git \
	vim \
	python-dev \
	python3-pip \
	build-essential \
	wget && \
	rm -rf /var/lib/apt/lists/*

# Install Anaconda for python 3.6
ADD Anaconda3-5.2.0-Linux-x86_64.sh /home/anaconda.sh
RUN /bin/bash /home/anaconda.sh -b -p /opt/conda && \
	ln -s /opt/conda/etc/profile.d/conda.sh /etc/profile.d/conda.sh && \
	rm /home/anaconda.sh

# Change sources for conda, add tsinghua sources and remove defaults
RUN conda config --add channels http://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free/ && \
	conda config --add channels http://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main/ && \
	conda config --add channels http://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/conda-forge/ && \
	conda config --add channels http://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/pytorch/ && \
	conda config --add channels http://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/pro/ && \
	conda config --remove channels defaults

# Change sources for pip3
RUN mkdir ~/.pip && \
	echo "[global]\nindex-url = http://mirrors.aliyun.com/pypi/simple/\n[install]\ntrusted-host = mirrors.aliyun.com" > ~/.pip/pip.conf

# Initialize workspace
RUN mkdir /workspace
WORKDIR /workspace

CMD ["/bin/bash"]

其中，Anaconda3-5.2.0-Linux-x86_64.sh 可以在 Anaconda Archive 中找到。

然后，我们可以通过下列指令制作镜像 my-nvidia/cuda:11.2。

cd ~/mkimage
sudo docker build -f Dockerfile -t my-nvidia/cuda:11.2 .

经过漫长的等待，我们可以看到一个 Successfully 提示消息，证明我们镜像打包成功。

构建完成后，我们可以通过下面这个指令进行容器的创建：

sudo docker run -it --gpus all --name cuda_11.2 my-nvidia/cuda:11.2 /bin/bash

进入容器之后，我们可以通过 nvidia-smi 、nvcc -V 、conda info 查看当前的显卡驱动、CUDA版本和 conda 源信息。
如果使用 conda 安装包的时候出现了conda Malformed version string ‘~’: invalid character(s)报错，可以使用下面的命令更新一下 conda。

conda upgrade -n base -c defaults --override-channels conda
conda update --all

我这里从 PyTorch 官网中下载了一个对应 CUDA 版本的 torch（我创建了一个名为 pytorch 的 conda 虚拟环境），可以看到在容器中 GPU 资源是可以正常被访问的。这样我们以后就可以随时切换 CUDA 版本了，是不是很方便？