首先,卸載掉舊版本torch torchvision:
pip uninstall torch torchvision
當然,也要卸載掉torch-geometric等依賴包。
下一步,安裝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
windows需要配置環境變量Path:
C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.2\lib\x64 C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.2\include C:\Program Files\NVIDIA GPU Computing Toolkit\CUDA\v10.2\extras\CUPTI\lib64 C:\ProgramData\NVIDIA Corporation\CUDA Samples\v10.2\bin\win64 C:\ProgramData\NVIDIA Corporation\CUDA Samples\v10.2\common\lib\x64
安裝完成后,cmd驗證:
nvcc -V
新版本torch 1.8.1分別支持cuda 10.2 和cuda 11.1. (https://pytorch.org/get-started/locally/)
支持cuda 10.2版本的torch安裝:
pip install torch==1.8.1+cu102 torchvision==0.9.0+cu102 torchaudio==0.8.0 -f https://download.pytorch.org/whl/torch_stable.html
支持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
安裝完成后,進入python驗證。
import torch
torch.__version__ 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 cpu
, cu101
, cu102
, or cu111
depending on your PyTorch installation.