torch-geometric

torch-geometric

技能概述

torch-geometric 技能提供 PyTorch Geometric (PyG) 图神经网络库的完整使用指南，涵盖从安装配置、核心概念理解到 GNN 模型构建、异构图处理和大图训练的全流程支持。

适用场景

图神经网络开发与训练

- 使用 PyG 构建和训练各类图神经网络模型，包括节点分类、图分类、链接预测等任务，支持 GCN、GAT、GraphSAGE、GIN 等主流 GNN 层。

异构图处理与分析

- 处理包含多种节点类型和边类型的异构图数据，适用于社交网络、知识图谱、推荐系统等复杂关系场景，提供 HeteroData 和 HGTConv 等专门工具。

大规模图数据训练

- 通过 NeighborLoader、ClusterLoader 等采样技术处理超出 GPU 内存的大图，支持分布式训练和 PyTorch Lightning 集成，适用于工业级图数据场景。

核心功能

完整的 GNN 构建体系

- 提供 60+ 预实现的图神经网络层（GCNConv、GATConv、SAGEConv、GINConv 等），支持自定义 MessagePassing 层，内置 GraphSAGE、GCN、GAT、GIN 等高级 API，满足从研究到生产的需求。

丰富的数据处理能力

- 内置 Planetoid、TUDataset、QM9 等标准数据集，支持 Data 和 HeteroData 图数据结构，提供 transforms 数据预处理管道，可从 CSV、NetworkX、scipy 等多种格式加载自定义图数据。

可扩展训练与解释

- 支持 NeighborLoader 邻居采样、ClusterGCN、GraphSAINT 等大图训练方案，提供 torch_geometric.explain 模块进行 GNN 预测解释（GNNExplainer、PGExplainer 等），兼容 torch.compile 和分布式训练。

常见问题

PyTorch Geometric 需要什么环境？

需要 Python 3.10+ 和 PyTorch 2.6+。推荐使用 uv pip install torch_geometric 安装核心库，扩展加速库（pyg-lib、torch-scatter 等）需根据你的 PyTorch 和 CUDA 版本从 https://data.pyg.org/whl 安装匹配的 wheel。

torch_geometric 和 NetworkX 有什么区别？

torch_geometric 专注于图神经网络训练，提供 GPU 加速的 GNN 层和大规模图训练支持；NetworkX 是通用图分析库，更适合图算法、可视化和小规模图计算。如果是深度学习任务用 PyG，经典图算法用 NetworkX。

如何使用 PyG 处理超出内存的大图？

使用 NeighborLoader 进行邻居采样：指定 num_neighbors 参数控制每层采样数量，batch_size 控制种子节点数。PyG 会构建包含采样邻居的子图进行训练，避免全图加载。对于链接预测任务使用 LinkNeighborLoader，多 GPU 场景可结合 torch.distributed DDP。