torch-geometric

PyTorch Geometric (PyG) - 图神经网络深度学习库

技能概述

PyTorch Geometric (PyG) 是基于 PyTorch 构建的图神经网络（GNN）开发库，专为处理图结构和几何深度学习任务设计。这个技能提供了 40+ 种图卷积层（如 GCN、GAT、GraphSAGE）、完整的消息传递框架、大规模图采样训练以及模型可解释性工具，适用于节点分类、图分类、链路预测、分子属性预测等多种场景。

适用场景

1. 图机器学习与深度学习研究

PyG 是图深度学习研究的核心工具，支持节点分类、图分类、链路预测等经典任务。无论是引文网络中的论文分类、社交网络中的社区发现，还是推荐系统中的关系建模，PyG 都提供了高效的数据加载、模型训练和评估流程。内置的 Planetoid、TU Dataset、QM9 等基准数据集让研究者能够快速验证算法。

2. 分子发现与化学属性预测

在药物研发和化学信息学领域，PyG 能够处理分子结构的图表示（原子作为节点，化学键作为边）。通过图神经网络预测分子的物理化学性质、药物活性或化学反应结果。PyG 支持处理 3D 分子几何结构，适用于几何深度学习任务，是计算化学和药物发现的重要工具。

3. 大规模异构图学习

对于包含多种节点和边类型的复杂图结构（如知识图谱、社交平台），PyG 提供了 HeteroData 类和 HeteroConv 层来处理异构图。配合 NeighborLoader 进行邻居采样，可以在单机上训练包含数百万节点的大规模图。支持多 GPU 并行训练，满足工业级应用需求。

核心功能

1. 丰富的图神经网络层与消息传递框架

PyG 内置 40+ 种图卷积层，包括 GCN（图卷积网络）、GAT（图注意力网络）、GraphSAGE 等。通过 MessagePassing 基类，用户可以轻松自定义消息传递层，实现特定的邻域聚合策略。框架自动处理 COO 格式的边索引、节点特征矩阵等图数据结构，支持稀疏矩阵运算和边权重、边属性等高级功能。

2. 高效的数据处理与训练流程

PyG 采用独特的批量处理机制，通过块对角邻接矩阵将多个图拼接成一个不连通的大图，无需填充即可高效计算。DataLoader 支持自动批处理和打乱，NeighborLoader 实现大规模图的邻居采样。对于超大数据集，可使用流式处理和分布式训练。数据增强变换（如添加自环、特征归一化、位置编码）可通过 Compose 链式调用。

3. 模型可解释性与可视化分析

PyG 提供了 GNNExplainer 等工具，可以分析模型预测的重要边和节点特征，帮助理解模型决策过程。支持节点掩码、边掩码等多种解释类型。配合可视化工具，可以直观展示图结构和注意力权重，对于科研和工业落地中的模型调试和信任度建立非常重要。

常见问题

PyTorch Geometric 怎么安装？

安装 PyG 非常简单，推荐使用 uv 或 pip：uv pip install torch_geometric。如需稀疏运算、聚类等加速功能，可额外安装 pyg_lib、torch_scatter、torch_sparse 等扩展包。需要注意的是，这些扩展包需要与您的 PyTorch 版本和 CUDA 版本匹配，官方提供了预编译 wheel 文件，可通过 -f 参数指定下载地址。

PyG 适合处理多大的图数据？

PyG 设计上支持从小规模到超大规模的图数据。对于内存可容纳的图（数万节点），可直接全图训练；对于大规模图（百万级节点），推荐使用 NeighborLoader 进行邻居采样训练。PyG 的批量处理机制对多个小图也非常高效。如果遇到内存瓶颈，可以减少采样邻居数、降低批次大小或使用流式数据加载。

PyG 和其他图神经网络库（如 DGL）如何选择？

PyG 与 DGL 都是优秀的图神经网络库。PyG 深度集成 PyTorch 生态，API 设计符合 PyTorch 习惯，对于已熟悉 PyTorch 的用户更友好。DGL 支持多种后端（MXNet、TensorFlow、PyTorch），在分布式训练方面更成熟。如果项目基于 PyTorch 且主要关注算法研究，推荐 PyG；如果需要跨框架或大规模工业部署，可考虑 DGL。