根据官方安装指南选择获取安装命令,例如:
conda install pytorch torchvision cudatoolkit=10.1 -c pytorch
使用国内镜像,比如清华源镜像,可以使用:
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main/
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free/
但是pytorch需要另外的源,添加如下源:
conda config --add channels https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud/pytorch
注意:PyTorch 官网给的安装命令需要去掉最后的 -c pytorch
,这样才能享受清华源的高速。安装的版本可能不正确,重新执行官网命令。
查看GPU信息
torch.cuda.is_available()
cuda是否可用;
torch.cuda.device_count()
返回gpu数量;
torch.cuda.get_device_name(0)
返回gpu名字,设备索引默认从0开始;
torch.cuda.current_device()
返回当前设备索引;