准备工作请看我的前两篇内容
Ubantu 20.04 安装 CUDA 、cuDNN并实现自由切换教程
1.检查安装所需条件
先到PyTorch官网查看安装所需条件,主要是看对于Python和CUDA版本的要求,比如我这里要求Python版本在3.8及以上,那么后面建立虚拟环境的时候就要注意这一点。
2.检查CUDA版本
我这里是11.8,所以在安装PyTorch页面选择CUDA 11.8。如果这里没有你的CUDA版本,可以查看我的上一篇文章。
nvcc --version
3.创建虚拟环境
建立一个名为 pytorch_test 的虚拟环境,以后所有的操作都在这个虚拟环境中。
conda create -n pytorch_test python=3.9
4.激活虚拟环境
可以看到,输入这个命令之后,最前面的环境名称变了。
conda activate pytorch_test
5.安装pytorch
直接复制官网给出的安装指令
conda install pytorch torchvision torchaudio pytorch-cuda=11.8 -c pytorch -c nvidia
6.检查是否成功
首先进入python环境
python
输入以下内容,返回True就是成功了。
import torch
torch.cuda.is_available()