torch cuda 配置

2021/4/14 18:25:33

本文主要是介绍torch cuda 配置,对大家解决编程问题具有一定的参考价值,需要的程序猿们随着小编来一起学习吧!

首先,卸载掉旧版本torch torchvision:

pip uninstall torch torchvision

当然,也要卸载掉torch-geometric等依赖包。

新版本torch 1.8.1分别支持cuda 10.2 和cuda 11.1. (https://pytorch.org/get-started/locally/)

支持cuda 10.2版本的torch安装:

pip install torch torchvision torchaudio

支持cuda 11.1版本的torch安装:

pip 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,以10.2版本为例。(https://developer.nvidia.com/cuda-10.2-download-archive?target_os=Linux&target_arch=x86_64&target_distro=Ubuntu&target_version=1804&target_type=runfilelocal)

wget https://developer.download.nvidia.com/compute/cuda/10.2/Prod/local_installers/cuda_10.2.89_440.33.01_linux.run

sh cuda_10.2.89_440.33.01_linux.run

安装完成后,进入python验证。

import torch
torch.cuda.is_available()

显示True则安装成功。

最后安装torch-geometric相关包。(https://github.com/rusty1s/pytorch_geometric/tree/9166ec13b6e005d9f1647f5f94efa058eef48202)

$ pip install torch-scatter -f https://pytorch-geometric.com/whl/torch-1.8.0+${CUDA}.html
$ pip install torch-sparse -f https://pytorch-geometric.com/whl/torch-1.8.0+${CUDA}.html
$ pip install torch-cluster -f https://pytorch-geometric.com/whl/torch-1.8.0+${CUDA}.html
$ pip install torch-spline-conv -f https://pytorch-geometric.com/whl/torch-1.8.0+${CUDA}.html
$ pip install torch-geometric

where ${CUDA} should be replaced by either cpucu101cu102, or cu111 depending on your PyTorch installation.

 



这篇关于torch cuda 配置的文章就介绍到这儿,希望我们推荐的文章对大家有所帮助,也希望大家多多支持为之网!


扫一扫关注最新编程教程