torch_geometric 是一个专门用于图神经网络(Graph Neural Networks, GNN) 的库,称为 PyTorch Geometric(简称 PyG),它并不包含在 PyTorch 的默认安装中,因此你需要单独安装它。

注意:此教程仅适用于CUDA和pytorch已经成功安装的情况。 需要根据他们的版本来确定PyG的安装版本。

1. 使用 Conda 安装

使用 Conda 安装与 CUDA 11.8 版本的 PyTorch 兼容的 torch_geometric,此时已经安装了 CUDA 11.8 版本、PyTorch。在虚拟环境中,使用以下命令安装:

conda install pyg -c pyg -c conda-forge

上述命令会从 pyg 和 conda-forge 频道安装 PyTorch Geometric 及其依赖,Conda 会根据你当前安装的 PyTorch 版本(包括 CUDA 11.8)自动选择兼容的版本。

上述命令太慢了,查看已安装torch版本2.0.0+cu118,CUDA版本11.8,利用以下pip命令安装:

pip install torch-geometric torch-scatter torch-sparse torch-cluster torch-spline-conv -f https://data.pyg.org/whl/torch-2.0.0+cu118.html

出现以下结果,即为安装完成。
在这里插入图片描述

2. 验证安装

完成安装后,可以使用以下 Python 代码来验证 torch_geometric 是否安装成功并兼容:

import torch
from torch_geometric.data import Data
from torch_geometric.nn import GCNConv

print('PyTorch Geometric installed and imported successfully')

成功安装后,会得到以下结果:
在这里插入图片描述
没有报错,说明 torch_geometric 安装成功并与当前的 PyTorch 版本(CUDA 11.8)兼容。

Logo

有“AI”的1024 = 2048,欢迎大家加入2048 AI社区

更多推荐