用途:PyTorch Geometric(PyG)是构建在 PyTorch 之上的图神经网络(GNN)库,用于图和不规则结构上的深度学习,涵盖节点分类、图分类、链接预测、分子性质预测、社交网络与引文网络分析、3D 几何数据(点云、网格、分子结构)、异构图(如知识图谱)以及大规模图学习。
主要能力:提供 40 多个卷积层,常用包括 GCNConv、GATConv、SAGEConv;支持基于 MessagePassing 的自定义消息传递层;内置 Planetoid、TUDataset、QM9、Reddit 等基准数据集;支持 DataLoader 小批量处理(块对角邻接矩阵拼接,无需填充)、NeighborLoader 邻居采样、多 GPU 训练;提供 HeteroData 与 to_hetero、HeteroConv 处理异构图;提供 NormalizeFeatures、AddSelfLoops、KNNGraph 等变换;提供 GNNExplainer 等可解释性工具与池化操作。
使用方式:安装命令为 uv pip install torch_geometric;如需稀疏操作、聚类等附加依赖,使用 uv pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-${TORCH}+${CUDA}.html。图数据用 torch_geometric.data.Data 表示,关键属性包括 data.x(节点特征)、data.edge_index(COO 格式 [2, num_edges] 连接关系)、data.edge_attr、data.y、data.pos,可扩展自定义属性。可从 CSV 加载节点与边构建图,或继承 InMemoryDataset/Dataset 创建自定义数据集。训练流程与 PyTorch 一致:定义模型、优化器,前向传播后计算损失并反向更新。
适用场景:图机器学习、药物发现与化学性质预测、社区检测与影响力预测、论文分类与推荐系统、点云与网格处理、知识图谱等。
限制:邻居采样通常不宜超过 2-3 跳;采样输出子图是有向的且节点索引会重排,损失只应使用种子节点预测;大规模图需借助采样或分布式训练。资料未提供具体收费、价格与平台支持信息。