1、Nvidia驱动安装参考https://www.cnblogs.com/a5idc/p/14003422.html
2、通过nvidia-smi命令查看cuda版本号
3、下载cuDNN
NVIDIA官网下载,根据平台选择,对应的版本,本次选择UBUNTU LINUX X86_64
是一个tar.gz文件,解压出来包括头文件和库,创建目录/usr/local/cuda/include和/usr/local/cuda/lib64,将加压出来的所有头文件和所有库分别拷贝到上述两个目录中。
4、安装Pytoch
Pytoch官网根据平台版本选择对的选项,
linux+pip3+python3+CUDA11.1生成的命令如下:
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是11.2版本,也是对应的torch也是1.8.1+cu111
5、检验torch安装是否ok
import torch
v_torch = torch.__version__
v_cuda = torch.version.cuda
v_cudnn = torch.backends.cudnn.version()
cuda_device_name = torch.cuda.get_device_name(0)
print("torch version: ", v_torch)
print("cuda version: ", v_cuda)
print("cudnn version: ", v_cudnn)
print("device name: ", cuda_device_name)
输出
torch version: 1.8.1+cu111
cuda version: 11.1
cudnn version: 8005
device name: GeForce GTX ...