PyTorch Geometric：图神经网络开发，一个库就够了

PyG 在 GitHub 上拿了 23.8K Star，是图神经网络领域最主流的 PyTorch 扩展库。

做 GNN 的人很难绕开它。学术圈发论文还是工业界搞分子预测、推荐系统、知识图谱，PyG 给了一套统一 API，10 到 20 行代码就能跑起一个图神经网络。

1、PyG 做了什么

在 PyTorch 上封装消息传递机制，让你不用从零实现 GCN、GAT、GraphSAGE 这些经典架构。图数据跟表格数据不一样。社交网络、分子结构、引文图谱天然就是图，PyG 把这层复杂性消化掉了。

几个核心能力：

• MessagePassing 是所有 GNN 层的基类，继承它实现 message() 和 forward() 就能定义自己的图卷积算子
• 自带 mini-batch loader，单张大图和小图批量训练都支持
• 通过采样方案处理百万节点级别的大规模图
• 多 GPU 训练和 torch.compile 加速跟上了 PyTorch 生态节奏

2、上手什么样

加载 Cora 引文数据集，定义一个两层 GCN，代码量少到让人意外：

from torch_geometric.nn import GCNConv
from torch_geometric.datasets import Planetoid

dataset = Planetoid(root='.', name='Cora')

class GCN(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = GCNConv(in_channels, hidden_channels)
        self.conv2 = GCNConv(hidden_channels, out_channels)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index).relu()
        x = self.conv2(x, edge_index)
        return x

PyTorch 用户不会有陌生感。张量为中心的设计，训练循环跟标准 PyTorch 一样：optimizer.zero_grad()、loss.backward()、optimizer.step()，三步不变。

自定义层也直接。想实现点云处理里的 EdgeConv，继承 MessagePassing，聚合方式指定 max，在 message() 里写消息计算逻辑就完了。图上稀疏运算、邻居聚合这些底层细节 PyG 已经封装好。

3、模型覆盖有多全

PyG 支持的 GNN 架构数量在同类库里排前列：

GNN 层：GCNConv、GATConv、GATv2Conv、SAGEConv、GINConv、TransformerConv、EdgeConv、PointNetConv 等超过 40 种。从谱域图卷积到图 Transformer，每层配有论文引用和示例代码。

池化层：Top-K Pooling、DiffPool、SAG Pooling、MinCut Pooling、Edge Pooling 等十几种。图分类中必不可少的层级表示学习，主流聚合策略都有覆盖。

完整模型：SchNet、DimeNet（分子建模）、Node2Vec（图嵌入）、Deep Graph Infomax（无监督学习）、GNNExplainer（可解释性）、GraphUNet 等，包含了预处理、跳跃连接、图粗化等完整设计。

4、四层架构

• Engine 层：基于 PyTorch，支持 torch.compile 和 TorchScript，配合 pyg-lib 做稀疏数据高效运算
• Storage 层：数据加载、转换和预处理，处理异构图和大规模图采样
• Operators 层：消息传递、聚合函数、归一化等构建块，自由组合
• Models 层：开箱即用的预定义模型和基准示例

日常使用 Models 层就够了，需要定制时下沉到 Operators 甚至直接改 MessagePassing。

5、适用方向

• GNN 方向的研究生和研究员：新想法验证快，baseline 现成
• 分子和材料科学：QM9、MD17 等数据集内置支持，SchNet、DimeNet 等专用模型直接可用
• 推荐系统和社交网络：GraphSAGE 风格的大规模图表征学习
• 3D 点云处理：PointNet++、EdgeConv、PointTransformer 等点云专用架构

文档和示例完备，从入门教程到 OGB 竞赛级代码都有。Slack 社区活跃。

dgeConv、PointTransformer 等点云专用架构

文档和示例完备，从入门教程到 OGB 竞赛级代码都有。Slack 社区活跃。