PyTorch Geometric — 面向图神经网络的主流开源工具库

PyG (PyTorch Geometric) 是一个基于 PyTorch 构建的库，用于轻松编写和训练图神经网络（GNN），适用于各种与结构化数据相关的应用。

它汇集了来自多篇已发表论文的方法，涵盖图及其他非规则结构上的深度学习（即几何深度学习）。
此外，它还提供了易于使用的 mini-batch 加载器（可用于处理大量小图和单个巨型图）、多 GPU 支持、torch.compile 支持、DataPipe 支持、大量常用基准数据集（提供简单接口以便创建自己的数据集），以及实用的数据变换工具，既适用于任意图上的学习，也适用于 3D 网格或点云。

[点击此处加入我们的 Slack 社区！][slack-url]

• 库亮点
• 新用户快速上手
• 架构概览
• 已实现的 GNN 模型
• 安装

库亮点

无论你是机器学习研究者，还是首次使用机器学习工具包的用户，以下理由都值得你尝试用 PyG 进行图结构数据的机器学习。

• 易用且统一的 API：
  只需 10 到 20 行代码即可开始训练 GNN 模型（参见下一节的快速上手）。
  PyG 堪称 PyTorch-on-the-rocks：它采用以张量为中心的 API，并保持与原生 PyTorch 相近的设计原则。
  如果你已熟悉 PyTorch，使用 PyG 将会非常轻松。
• 全面且维护良好的 GNN 模型：
  大多数最先进的图神经网络架构已由库开发者或研究论文作者实现，可直接应用。
• 出色的灵活性：
  现有 PyG 模型可以方便地扩展，以支撑你自己的 GNN 研究。
  得益于易用的消息传递 API 以及丰富的算子与工具函数，修改现有模型或创建新架构都非常简单。
• 大规模真实场景 GNN 模型：
  我们关注 GNN 在具有挑战性的真实场景中的应用需求，并支持在多种类型的图上学习，包括但不限于：可扩展到数百万节点图的可扩展 GNN；用于随时间进行节点预测的动态 GNN；具有多种节点类型和边类型的异构图 GNN。

新用户快速上手

在本快速上手部分，我们将展示如何仅用几行代码轻松创建并训练一个 GNN 模型。

训练你自己的 GNN 模型

作为 PyG 的初次体验，我们实现一个用于引文图中论文分类的 GNN 训练流程。
为此，我们加载 Cora 数据集，并使用预定义的 GCNConv 创建一个简单的两层 GCN 模型：

import torch
from torch import Tensor
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: Tensor, edge_index: Tensor) -> Tensor:
        # x: 形状为 [num_nodes, in_channels] 的节点特征矩阵
        # edge_index: 形状为 [2, num_edges] 的图连接矩阵
        x = self.conv1(x, edge_index).relu()
        x = self.conv2(x, edge_index)
        return x

model = GCN(dataset.num_features, 16, dataset.num_classes)

import torch.nn.functional as F

data = dataset[0]
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

for epoch in range(200):
    pred = model(data.x, data.edge_index)
    loss = F.cross_entropy(pred[data.train_mask], data.y[data.train_mask])

    # 反向传播
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

有关评估最终模型性能的更多信息，请参见对应的示例。

创建你自己的 GNN 层

除了轻松应用现有 GNN 之外，PyG 还让实现自定义图神经网络变得简单（配套教程见此处）。
例如，实现 Wang 等人 提出的边卷积层只需以下代码：

$$x_i^{\prime} ~ = ~ \max_{j \in \mathcal{N}(i)} ~ \textrm{MLP}_{\theta} \left( [ ~ x_i, ~ x_j - x_i ~ ] \right)$$

import torch
from torch import Tensor
from torch.nn import Sequential, Linear, ReLU
from torch_geometric.nn import MessagePassing

class EdgeConv(MessagePassing):
    def __init__(self, in_channels, out_channels):
        super().__init__(aggr="max")  # "Max" 聚合。
        self.mlp = Sequential(
            Linear(2 * in_channels, out_channels),
            ReLU(),
            Linear(out_channels, out_channels),
        )

    def forward(self, x: Tensor, edge_index: Tensor) -> Tensor:
        # x: 形状为 [num_nodes, in_channels] 的节点特征矩阵
        # edge_index: 形状为 [2, num_edges] 的图连接矩阵
        return self.propagate(edge_index, x=x)  # 形状为 [num_nodes, out_channels]

    def message(self, x_j: Tensor, x_i: Tensor) -> Tensor:
        # x_j: 形状为 [num_edges, in_channels] 的源节点特征
        # x_i: 形状为 [num_edges, in_channels] 的目标节点特征
        edge_features = torch.cat([x_i, x_j - x_i], dim=-1)
        return self.mlp(edge_features)  # 形状为 [num_edges, out_channels]

架构概览

PyG 提供了一个多层框架，使用户既能在低层也能在高层构建图神经网络解决方案。
它由以下组件构成：

• PyG 引擎 依托强大的 PyTorch 深度学习框架，完整支持 torch.compile 和 TorchScript，并加入了用于稀疏数据运算的高效 CPU/CUDA 库，例如 pyg-lib。
• PyG 存储 负责数据处理、变换和加载流程。它能够处理并加工大规模图数据集，并为异构图提供有效解决方案。它还提供了多种采样方案，使得在大规模图上训练 GNN 成为可能。
• PyG 算子 汇集了实现图神经网络所需的核心功能。PyG 支持重要的 GNN 构建模块，这些模块可以组合并应用到 GNN 模型的各个部分，确保 GNN 设计具有充分的灵活性。
• 最后，PyG 提供了丰富的 GNN 模型集，以及展示 GNN 模型在标准图基准上表现的示例。得益于其灵活性，用户可以轻松构建和修改自定义 GNN 模型以满足特定需求。

已实现的 GNN 模型

我们按类别列出当前支持的 PyG 模型、层和算子：

GNN 层：
所有图神经网络层都通过 nn.MessagePassing 接口实现。
GNN 层定义了如何执行消息传递，即通过设计不同的消息、聚合和更新函数（定义见此处）。
这些 GNN 层可以堆叠在一起，构建图神经网络模型。

• GCNConv，来自 Kipf 和 Welling：Semi-Supervised Classification with Graph Convolutional Networks（ICLR 2017）[示例]
• ChebConv，来自 Defferrard 等人：Convolutional Neural Networks on Graphs with Fast Localized Spectral Filtering（NIPS 2016）[示例]
• GATConv，来自 Veličković 等人：Graph Attention Networks（ICLR 2018）[示例]

池化层：
图池化层将图中（或子图中）的一组节点的向量表示合并为单个向量表示，以汇总其节点属性。
它通常应用于图级任务，这类任务需要将节点特征组合成单一的图表示。

• Top-K 池化，来自 Gao 和 Ji：Graph U-Nets（ICML 2019）、Cangea 等人：Towards Sparse Hierarchical Graph Classifiers（NeurIPS-W 2018）以及 Knyazev 等人：Understanding Attention and Generalization in Graph Neural Networks（ICLR-W 2019）[示例]
• DiffPool，来自 Ying 等人：Hierarchical Graph Representation Learning with Differentiable Pooling（NeurIPS 2018）[示例]