安装PyTorch和PyTorch Geometric

2021-07-02
2 min read

1. 安装PyTorch

官网选择操作系统、下载方式和对应的GPU配置:

如果需要之前的版本可以点击下方链接选择之前的版本:

需要注意对应的CUDA版本号。如何安装、切换多版本CUDA可以参考此前的博客。可通过以下命令查看当前CUDA版本:

cat  /usr/local/cuda/version.txt

下载完成后可以通过以下命令检查torch版本及对应的CUDA版本:

python -c "import torch; print(torch.__version__)"
python -c "import torch; print(torch.version.cuda)"

可以通过以下命令查看GPU是否可用:

python
>>> import torch
>>> torch.cuda.is_available() 	# GPU是否可用
>>> torch.cuda.device_count()	# GPU数量
>>> torch.cuda.current_device()	# 当前GPU
>>> exit()

例如:

注意,GPU devices从0开始编号。

2. 安装PyTorch Geometric

2.1 快速安装

根据官网,如果PyTorch版本$\ge1.8.0$,可以快速下载:

2.2 自定义安装

自定义下载需要根据当前的PyTorch版本和CUDA版本下载相关的依赖,下载命令如下:

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

其中,${TORCH}替换为当前环境下的PyTorch版本,目前支持1.4.01.5.01.6.01.7.01.7.11.8.01.8.1、和1.9.0${CUDA}替换为指定的CUDA版本,目前支持cpucu92cu101cu102cu110cu111

例如对于PyTorch 1.8.0/1.8.1和CUDA 10.2:

pip install torch-scatter -f https://pytorch-geometric.com/whl/torch-1.8.0+cu102.html
pip install torch-sparse -f https://pytorch-geometric.com/whl/torch-1.8.0+cu102.html
pip install torch-geometric

2.3 版本依赖

使用自定义安装时,依然可能会出现安装失败的问题,因为pytorch geometric几个相关库之间有比较强的依赖关系,建议是在自定义安装的基础上指定对应库的版本,例如对于pytorch==1.4.0cuda==10.1

pip install torch-scatter==2.0.4 -f https://pytorch-geometric.com/whl/torch-1.4.0+cu101.html
pip install torch-cluster==1.5.4 -f https://pytorch-geometric.com/whl/torch-1.4.0+cu101.html
pip install torch-sparse==0.6.1 -f https://pytorch-geometric.com/whl/torch-1.4.0+cu101.html
pip install torch-spline-conv==1.2.0 -f https://pytorch-geometric.com/whl/torch-1.4.0+cu101.html
pip install torch-geometric==1.6.0 -f https://pytorch-geometric.com/whl/torch-1.4.0+cu101.html

每个库的版本信息可以在GitHub上的release页找到,例如:

几个相关库的GitHub地址如下: