在RTX 3090上安装Pytorch-1.6

近来因为复现别人的工作需要，必须使用1.6版本的Pytorch，但是Pytorch官方说明1.6版本最高只支持CUDA10.2，而实验室最近服务器升级了RTX 3090，至少要求CUDA11，因此两者无法兼容。

经过一通折腾，最后采用在服务器上安装nvidia-docker来解决这一需求。

初始时服务器的配置如下：

• 操作系统：Ubuntu 20.04
• Nvidia-driver：495.46
• CUDA：11.5
• cudnn：8.3.2

安装nvidia-docker

Installation Guide — NVIDIA Cloud Native Technologies documentation

curl https://get.docker.com | sh \
  && sudo systemctl --now enable docker
  
distribution=$(. /etc/os-release;echo $ID$VERSION_ID) \
      && curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg \
      && curl -s -L https://nvidia.github.io/libnvidia-container/$distribution/libnvidia-container.list | \
            sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \
            sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list

sudo apt-get update
sudo apt-get install -y nvidia-docker2
sudo systemctl restart docker

下载Pytorch镜像

Frameworks Support Matrix :: NVIDIA Deep Learning Frameworks Documentation

下载pytorch-1.6.0+CUDA11的镜像，在上述Nvidia官网找到20.06版本符合需求

PyTorch | NVIDIA NGC

docker pull nvcr.io/nvidia/pytorch:20.06-py3

这里在下载的时候由于镜像较大，服务器上的网络不稳定，因此采用在别的网络环境好的机器上下载然后上传到服务器的方法

# 机器A
docker pull nvcr.io/nvidia/pytorch:20.06-py3 # 下载镜像
docker images # 查看镜像ID
docker save ${IMAGE ID} > pytorch_image.tar # 打包镜像文件

# 将pytorch_image.tar文件从机器A传到服务器B

# 服务器B
docker load -i pytorch_image.tar # 加载镜像

运行Pytorch镜像

docker images # 查看镜像ID
docker run -it --gpus all --name ${CONTAINER NAME} -v 本地文件路径1：镜像中的文件路径1 -v 本地文件路径2：镜像中的文件路径2 --network host ${IMAGE ID} /bin/bash
# 其中${CONTAINER NAME}是自己给运行的container起的名字，由于需要代码数据模型等放入container中运行，因此-v后面两组路径分别是传递代码等内容，需要注意的是这里的两组路径相互可以修改，即在container对对应目录中的文件进行修改后，会在服务器中对应的本地路径上真实进行修改，因此可以通过这种方式保存container中运行的结果

# 此时可能会出现如下报错，但是不用管它，实际上是可以用的，看上去也不太影响性能
# WARNING: Detected NVIDIA NVIDIA GeForce RTX 3090 GPU, which is not yet supported in this version of the container
# ERROR: No supported GPU(s) detected to run this container

启动镜像进行容器后，首先安装tmux方便使用

apt-get install tmux
tmux new -s ${SESSION NAME} # 新建session从container进入tmux

# 修改代码启动程序等
# ...

# 退出时为了保证程序能够不中断，需要小心
# Ctrl-B + D # 退出tmux回到container
# Ctrl + P + Q # 从container临时退出，注意顺序不能错


# 现在就从container回到了服务器环境，可以进行其他操作
# 需要重新进入container查看实验结果时
docker ps -all # 查看container状态
docker attach ${CONTAINER NAME} # 从服务器重新进入container
tmux at -t ${SESSION NAME} # 重新进去tmux
# 这时就可以看到之前正在运行的程序，再次退出时同上次操作
# Ctrl-B + D # 退出tmux回到container
# Ctrl + P + Q # 从container返回服务器

# 当不再使用容器时，可以直接退出容器
exit
docker rm ${CONTAINER NAME} # 删除container