pytorch-gpu安装

发表于 分类于 阅读次数：
不同版本cuda下的安装pytorch-gpu的pip命令

nvcc -V查看cuda版本

CUDA 11.2

1
conda install pytorch==1.8.0 torchvision==0.9.0 torchaudio==0.8.0 cudatoolkit=11.2 -c pytorch -c conda-forge

CUDA 11.6

1
pip3 install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu116

测试：

1
2
3
4
5
6
import torch

torch.cuda.is_available()       # cuda是否可用
torch.cuda.current_device()     # 返回当前设备索引
torch.cuda.device_count()       # 返回GPU的数量
torch.cuda.get_device_name(0)   # 返回gpu名字，设备索引默认从0开始