跳到主要内容

简介与安装

简介​

pytorch是基于python的深度学习框架

(可选)安装CUDA​

虽说pytorch也支持纯CPU，但如果GPU支持CUDA的话，最好还是用CUDA来学习pytorch，计算速度会快很多。

参考nvidia官网和pytorch的版本来安装cuda。

安装完后验证

nvcc -V

(可选)安装cuDNN​

有了CUDA，还需要安装cuDNN。cuDNN是一个深度神经网络的加速库，实现了神经网络常见的操作。

安装​

  • windows: cuda 12.1
python -m pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121

验证​

import torch

print(f"Hello World, Hello PyTorch {torch.__version__}")
print(f"CUDA is available:{torch.cuda.is_available()}, version is {torch.version.cuda}")
print(f"device_name: {torch.cuda.get_device_name(0)}")

运行输出

Hello World, Hello PyTorch 2.3.1+cu121
CUDA is available:True, version is 12.1
device_name: NVIDIA GeForce RTX 4070 Ti SUPER

参考​