PyTorch深度学习环境搭建

现在已经是2021年了，搭建PyTroch环境也是越来越简单了，网上很多教程其实都已经过时了。现在安装PyTorch其实步骤只需要包括两步，1）安装NVIDIA驱动；2）按照PyTorch官网的指示安装PyTorch。本文将以Ubuntu 20.04为例，讲解如何搭建PyTroch环境。

1. 安装NVIDIA驱动

安装NVIDIA驱动有很多方法，可以在官网下载run二进制文件，Ubuntu系统也可以下载驱动的deb包。但是这些方法都比较复杂，最近Ubuntu的官方源就已经加入了NVIDIA的驱动了，直接用apt安装即可，省时省力，方便维护。

先搜索一下可用的NVIDIA驱动安装包：

apt-cache search nvidia-driver

然后会列出所有满足条件的包（为了简洁，下面我省略了一部分输出结果）：

...
nvidia-driver-418-server - NVIDIA Server Driver metapackage
nvidia-driver-440-server - Transitional package for nvidia-driver-450-server
nvidia-driver-460-server - NVIDIA Server Driver metapackage
...

挑一个最新的安装即可，因为我是安装在服务器上面，所以选择了包nvidia-driver-460-server。安装成功后需要重启一下，重启之后运行一下nvidia-smi命令，如果没有报错并且成果返回了显卡信息则说明驱动已经安装成果了。

2. 按照PyTorch官网的指示安装PyTorch

PyTorch已经内置了CUDA和CuDNN，所以无需在服务器上安装这两个程序了。直接按照PyTorch官网的指令按照对应的版本就可以了。比如我需要装CUDA 11.1版本的PyTorch，命令如下：

pip3 install torch==1.8.1+cu111 torchvision==0.9.1+cu111 torchaudio==0.8.1 \
-f https://download.pytorch.org/whl/torch_stable.html

安装完成之后运行以下命令查看是否能使用CUDA进行计算：

python -c 'import torch; print(torch.cuda.is_available())'

如果返回的结果是True则说明已经安装成功。