用途:PyTorch Geometric(PyG)是基于 PyTorch 的图神经网络(GNN)库,用于构建和训练图神经网络,适用于节点分类、图分类、链接预测等任务。
主要能力:提供图数据结构 Data 与 HeteroData;60 多种 GNN 卷积层(如 GCNConv、GATConv、SAGEConv、GINConv、TransformerConv、EdgeConv、RGCNConv、HGTConv);内置数据集(Planetoid、TUDataset、OGB、QM7/QM9 等);数据变换(NormalizeFeatures、ToUndirected、AddSelfLoops、KNNGraph、RandomJitter 等);高层模型 API(GraphSAGE、GCN、GAT、GIN);通过 MessagePassing 自定义层;大规模图邻居采样(NeighborLoader、LinkNeighborLoader、ClusterLoader、GraphSAINTSampler);异构图支持(to_hetero、HeteroConv、HGTConv);可解释性(GNNExplainer、PGExplainer、CaptumExplainer、AttentionExplainer)。
使用方式:安装命令为 uv add torch_geometric(或 uv pip install torch_geometric),需先安装 PyTorch;可选安装 pyg-lib、torch-scatter、torch-sparse、torch-cluster 以加速运算。核心用法是构造 Data 对象(x、edge_index、edge_attr、y、pos),其中 edge_index 必须为 [2, num_edges] 的 COO 格式张量。卷积层不内置激活函数,需自行添加。图分类任务使用 DataLoader 与 global_mean_pool 等池化操作。大图训练使用 NeighborLoader 进行邻居采样,num_neighbors 长度应与 GNN 层数一致。
适用场景:学术研究与工程中的图结构数据建模,包括引文网络、分子图、点云、知识图谱、推荐系统等。
限制:采样超过 2-3 跳通常不可行;异构图二部卷积层不应启用自环,应使用跳跃连接;使用 -1 输入通道的惰性初始化模型需先进行一次 torch.no_grad() 前向传播。具体收费、价格与平台信息暂未提供。