RTX3090深度学习环境配置（PyTorch1.8）

更新于2021年3月6日，昨天PyTorch官方已经发布了1.8版本，正式支持CUDA11。

简介

最近要在3090卡的服务器上进行PyTorch实验，所以踩了不少坑进行环境配置，本文就简单介绍一下我成功配置的流程，由于实验中只使用了PyTorch，所以其他框架的我这里没有测试，也不多做说明，此外本文设计的内容都是在Ubuntu系统下进行的。

显卡驱动

显卡驱动直接到官方下载runfile文件进行安装即可，这方面有不少教程，我这里不多赘述，有兴趣的可以参考我的Ubuntu教程。

成功安装驱动之后，命令行输入nvidia-smi应该会出现所有本机的显卡信息，如6卡信息如下图，这代表显卡驱动安装成功，右上角的CUDA Version字样则表示该驱动支持最高CUDA版本为11.1，这里尤其需要注意的是，30系列显卡安装的CUDA必须为11，至于具体为11.0还是11.1，官方建议是11.1，因为11.0会出现很多异常。

PyTorch环境

系统CUDA我这边预装的是11.1（下图命令查看），具体安装也只要到官网下载runfile文件就可以了，我这边就略过，因为这个系统CUDA可装可不装，也可以选择虚拟环境里通过cudatoolkit获得的cuda。

首先，创建一个环境，命令如下。

conda create -n torch python=3.8 -y
conda activate torch

接着，就是安装PyTorch了，需要注意的是，目前PyTorch发布了1.8版本，已经稳定支持CUDA11.1，因此下面我们使用目前PyTorch最高支持的CUDA11.1来配置。这里使用的是新推出的LTS版本，由于源在国外，因此可能需要比较长时间的耐心等待。

conda install pytorch torchvision cudatoolkit=11.1 -c pytorch-lts -c nvidia

当然，如果你的系统已经安装了CUDA11.1，那么可以去掉其中的cudatoolkit包的安装。

如果你的系统正确安装了CUDA11.1和对应的CUDNN那么，此时就可以正常运行了。如果没有安装CUDNN，需要自己到官网下载对应版本，比如这里对应于cuda11.1，我下载的是cudnn8.0.5，cudnn压缩包解压后得到的cuda文件夹内的lib64目录下的所有文件拷贝到当前cuda环境的lib文件夹下，以我的test环境为例，命令如下。

sudo cp cuda/include/cudnn.h /usr/local/cuda/include
sudo cp cuda/lib64/libcudnn* /usr/local/cuda/lib64
sudo chmod a+r /usr/local/cuda/include/cudnn.h
sudo chmod a+r /usr/local/cuda/lib64/libcudnn*

最后，测试是否成功配置环境，命令为python -c "import torch;print(torch.__version__);print(torch.cuda.is_available())"，结果如下，可以发现，测试是通过的，实际上我也去跑了实际的程序，速度只能说正常，显存占用相比于我同样配置在2080Ti上跑不知道为什么高了很多，可能因为太新的缘故，优化没做好吧，等待稳定版的更新。

总结

本文介绍目前（2022.3月）最新的在3090显卡机器上高效配置PyTorch环境的方式（TensorFlow的话可以搜索其他教程），如果我的文章对你有所帮助，欢迎一键三连。