torch-geometric

PyTorch Geometric (PyG)

PyG 是基于 PyTorch 构建的图神经网络标准库。它提供图数据结构、60 多种 GNN 层实现、可扩展的迷你批次训练，以及对异构图的支持。

安装：uv add torch_geometric（或 uv pip install torch_geometric；需要 PyTorch）。可选：pyg-lib、torch-scatter、torch-sparse、torch-cluster，用于加速运算。

核心概念

图数据：Data 和 HeteroData

图存储在 Data 对象中。关键属性：

from torch_geometric.data import Data

data = Data(
    x=node_features,          # [num_nodes, num_node_features]
    edge_index=edge_index,     # [2, num_edges] — COO format, dtype=torch.long
    edge_attr=edge_features,   # [num_edges, num_edge_features]
    y=labels,                  # node-level [num_nodes, *] or graph-level [1, *]
    pos=positions,             # [num_nodes, num_dimensions] (for point clouds/spatial)
)

edge_index 的格式至关重要：它是一个形状为 [2, num_edges] 的张量，其中 edge_index[0] = 源节点，edge_index[1] = 目标节点。它不是元组列表。如果边对以行为单位排列，请转置并调用 .contiguous()：

# If edges are [[src1, dst1], [src2, dst2], ...] — transpose first:
edge_index = edge_pairs.t().contiguous()

对于无向图，需要包含两个方向：边 (0,1) 在 edge_index 中需要同时包含 [0,1] 和 [1,0]。

对于异构图，请使用 HeteroData——参见下文的异构图部分。

数据集

PyG 附带许多标准数据集，可自动下载和预处理：

from torch_geometric.datasets import Planetoid, TUDataset

# Single-graph node classification (Cora, Citeseer, Pubmed)
dataset = Planetoid(root='./data', name='Cora')
data = dataset[0]  # single graph with train/val/test masks

# Multi-graph classification (ENZYMES, MUTAG, IMDB-BINARY, etc.)
dataset = TUDataset(root='./data', name='ENZYMES')
# dataset[0], dataset[1], ... are individual graphs

按任务划分的常见数据集：

• 节点分类：Planetoid（Cora/Citeseer/Pubmed）、OGB（ogbn-arxiv、ogbn-products、ogbn-mag）
• 图分类：TUDataset（MUTAG、ENZYMES、PROTEINS、IMDB-BINARY）、OGB（ogbg-molhiv）
• 链接预测：OGB（ogbl-collab、ogbl-citation2）
• 分子：QM7、QM9、MoleculeNet
• 点云/网格：ShapeNet、ModelNet10/40、FAUST

变换

变换用于预处理或增强图数据，类似于 torchvision 变换：

import torch_geometric.transforms as T

# Common transforms
T.NormalizeFeatures()    # Row-normalize node features to sum to 1
T.ToUndirected()         # Add reverse edges to make graph undirected
T.AddSelfLoops()         # Add self-loop edges
T.KNNGraph(k=6)          # Build k-NN graph from point cloud positions
T.RandomJitter(0.01)     # Random noise augmentation on positions
T.Compose([...])         # Chain multiple transforms

# Apply as pre_transform (once, saved to disk) or transform (every access)
dataset = ShapeNet(root='./data', pre_transform=T.KNNGraph(k=6),
                   transform=T.RandomJitter(0.01))

构建 GNN 模型

快速入门：使用内置层

构建 GNN 最快的方法是从 torch_geometric.nn 堆叠卷积层：

import torch
import torch.nn.functional as F
from torch_geometric.nn import GCNConv

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 = F.dropout(x, p=0.5, training=self.training)
        x = self.conv2(x, edge_index)
        return x

重要提示：PyG 卷积层不包含激活函数——请在每一层之后自行应用。这是为了提高灵活性而特意设计的。

选择卷积层

根据任务和图结构选择：

层	最适用于	核心思路
GCNConv	同构图、半监督节点分类	受谱方法启发，按度归一化聚合
GATConv / GATv2Conv	邻居的重要性各不相同时	基于注意力加权的消息
SAGEConv	大型图、归纳式场景	适合采样的可学习聚合
GINConv	图分类、最大化表达能力	与 WL 测试同样强大
TransformerConv	丰富的边特征、复杂交互	使用边特征的多头注意力
EdgeConv	点云、动态图	对边特征 (x_i, x_j - x_i) 使用 MLP
RGCNConv	关系类型较多的异构图	针对关系类型的权重矩阵
HGTConv	异构图	按类型区分的注意力

所有卷积层至少接受 (x, edge_index)。许多层还接受用于表示边特征的 edge_attr。

惰性初始化

输入通道数使用 -1，可让 PyG 自动推断维度——对于异构模型尤其有用：

conv = SAGEConv((-1, -1), 64)  # Input dims inferred on first forward pass
# Initialize lazy modules:
with torch.no_grad():
    out = model(data.x, data.edge_index)

高层模型 API

对于常见架构，PyG 提供了可直接使用的模型类：

from torch_geometric.nn import GraphSAGE, GCN, GAT, GIN

model = GraphSAGE(
    in_channels=dataset.num_features,
    hidden_channels=64,
    out_channels=dataset.num_classes,
    num_layers=2,
)

通过 MessagePassing 自定义层

若要实现新的 GNN 层，可继承 MessagePassing。其框架如下：

1. propagate() 协调消息传递
2. message() 定义沿每条边传递的信息（phi 函数）
3. aggregate() 汇总每个节点收到的消息（sum/mean/max）
4. update() 转换聚合后的结果（gamma 函数）

from torch_geometric.nn import MessagePassing
from torch_geometric.utils import add_self_loops, degree

class MyConv(MessagePassing):
    def __init__(self, in_channels, out_channels):
        super().__init__(aggr='add')  # "add", "mean", or "max"
        self.lin = torch.nn.Linear(in_channels, out_channels)

    def forward(self, x, edge_index):
        # Pre-processing before message passing
        x = self.lin(x)
        # Start message passing
        return self.propagate(edge_index, x=x)

    def message(self, x_j):
        # x_j: features of source nodes for each edge [num_edges, features]
        # The _j suffix auto-indexes source nodes, _i indexes target nodes
        return x_j

_i / _j 约定：传递给 propagate() 的任意张量，都可以在 message() 的签名中追加 _i（目标/中心节点）或 _j（源/邻居节点）来自动索引。因此，如果向 propagate 传递 x=...，就可以在 message() 中访问 x_i 和 x_j。

阅读 references/message_passing.md，查看完整的 GCN 和 EdgeConv 实现示例。

特定任务模式

节点分类

# Full-batch training on a single graph (e.g., Cora)
model.train()
for epoch in range(200):
    optimizer.zero_grad()
    out = model(data.x, data.edge_index)
    loss = F.cross_entropy(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()

# Evaluation
model.eval()
pred = model(data.x, data.edge_index).argmax(dim=1)
acc = (pred[data.test_mask] == data.y[data.test_mask]).float().mean()

图分类

多个图：使用 DataLoader 进行小批量处理，并通过全局池化获取图级表示：

from torch_geometric.loader import DataLoader
from torch_geometric.nn import GCNConv, global_mean_pool

loader = DataLoader(dataset, batch_size=32, shuffle=True)

class GraphClassifier(torch.nn.Module):
    def __init__(self, in_ch, hidden_ch, out_ch):
        super().__init__()
        self.conv1 = GCNConv(in_ch, hidden_ch)
        self.conv2 = GCNConv(hidden_ch, hidden_ch)
        self.lin = torch.nn.Linear(hidden_ch, out_ch)

    def forward(self, x, edge_index, batch):
        x = self.conv1(x, edge_index).relu()
        x = self.conv2(x, edge_index).relu()
        x = global_mean_pool(x, batch)  # [num_graphs_in_batch, hidden_ch]
        return self.lin(x)

# Training loop
for data in loader:
    out = model(data.x, data.edge_index, data.batch)
    loss = F.cross_entropy(out, data.y)

PyG 的 DataLoader 通过创建块对角邻接矩阵来批处理多个图。batch 张量将每个节点映射到其图索引。池化操作（global_mean_pool、global_max_pool、global_add_pool）利用该张量按图进行聚合。

链路预测

将边拆分为训练集、验证集和测试集，并使用负采样：

from torch_geometric.transforms import RandomLinkSplit

transform = RandomLinkSplit(
    num_val=0.1,
    num_test=0.1,
    is_undirected=True,
    add_negative_train_samples=False,
)
train_data, val_data, test_data = transform(data)

# Encode nodes, then score edges
z = model.encode(train_data.x, train_data.edge_index)
# Positive edges
pos_score = (z[train_data.edge_label_index[0]] * z[train_data.edge_label_index[1]]).sum(dim=1)

阅读 references/link_prediction.md，获取完整的链路预测指南：GAE/VGAE 自编码器、完整训练循环、适用于大图的 LinkNeighborLoader、异构链路预测和评估指标。

扩展到大图

对于无法放入 GPU 内存的图，使用 NeighborLoader 进行邻居采样：

from torch_geometric.loader import NeighborLoader

train_loader = NeighborLoader(
    data,
    num_neighbors=[15, 10],     # Sample 15 neighbors in hop 1, 10 in hop 2
    batch_size=128,              # Number of seed nodes per batch
    input_nodes=data.train_mask, # Which nodes to sample from
    shuffle=True,
)

for batch in train_loader:
    batch = batch.to(device)
    out = model(batch.x, batch.edge_index)
    # Only use first batch_size nodes for loss (these are the seed nodes)
    loss = F.cross_entropy(out[:batch.batch_size], batch.y[:batch.batch_size])

关于 NeighborLoader 的要点：

• num_neighbors 列表的长度应与 GNN 深度（消息传递层数）相匹配
• 种子节点始终是输出中最前面的 batch.batch_size 个节点
• batch.n_id 将重新标记后的索引映射回原始节点 ID
• 适用于 Data 和 HeteroData
• 对于链路预测，请改用 LinkNeighborLoader
• 通常无法采样超过 2-3 跳（会呈指数级增长）

其他扩展方案：ClusterLoader（ClusterGCN）、GraphSAINTSampler、ShaDowKHopSampler。有关多 GPU 训练、DDP、PyTorch Lightning 集成和 torch.compile 支持，请阅读 references/scaling.md。

异构图

对于包含多种节点和边类型的图（社交网络、知识图谱、推荐）：

from torch_geometric.data import HeteroData

data = HeteroData()

# Node features — indexed by node type string
data['user'].x = torch.randn(1000, 64)
data['movie'].x = torch.randn(500, 128)

# Edge indices — indexed by (src_type, edge_type, dst_type) triplet
data['user', 'rates', 'movie'].edge_index = torch.randint(0, 500, (2, 3000))
data['user', 'follows', 'user'].edge_index = torch.randint(0, 1000, (2, 5000))

# Access convenience dicts
data.x_dict        # {'user': tensor, 'movie': tensor}
data.edge_index_dict  # {('user','rates','movie'): tensor, ...}
data.metadata()    # ([node_types], [edge_types])

构建异构 GNN 的三种方式

1. 使用 to_hetero() 自动转换 — 编写同构模型并自动转换：

from torch_geometric.nn import SAGEConv, to_hetero

class GNN(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = SAGEConv((-1, -1), hidden_channels)
        self.conv2 = SAGEConv((-1, -1), out_channels)

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

model = GNN(64, dataset.num_classes)
model = to_hetero(model, data.metadata(), aggr='sum')

# Now accepts dicts:
out = model(data.x_dict, data.edge_index_dict)

对于二部图输入通道（源和目标可能不同），使用 (-1, -1)。其余部分由延迟初始化处理。

2. 使用 HeteroConv 封装器 — 为每种边类型使用不同的卷积层：

from torch_geometric.nn import HeteroConv, GCNConv, SAGEConv, GATConv

conv = HeteroConv({
    ('paper', 'cites', 'paper'): GCNConv(-1, 64),
    ('author', 'writes', 'paper'): SAGEConv((-1, -1), 64),
    ('paper', 'rev_writes', 'author'): GATConv((-1, -1), 64, add_self_loops=False),
}, aggr='sum')

3. 使用原生异构算子，例如 HGTConv：

from torch_geometric.nn import HGTConv
conv = HGTConv(hidden_channels, hidden_channels, data.metadata(), num_heads=4)

异构图的重要注意事项：

• 使用 T.ToUndirected() 添加反向边类型，以支持双向消息流
• 在二部图卷积层（源类型和目标类型不同）中禁用 add_self_loops — 改用跳跃连接：conv(x, edge_index) + lin(x)
• 在 HeteroData 上使用 NeighborLoader 时，将 input_nodes 指定为 ('node_type', mask) 元组
• num_neighbors 可以是按边类型为键的字典，以实现精细控制

阅读 references/heterogeneous.md，获取完整示例，包括训练循环以及在异构图上使用 NeighborLoader。

自定义数据集

将自己的数据加载到 PyG 中：

• 快速（无需类）：直接创建 Data 对象，并将列表传递给 DataLoader
• 可复用（可放入 RAM）：继承 InMemoryDataset — 重写 raw_file_names、processed_file_names、download()、process()
• 大型（基于磁盘）：继承 Dataset — 还需重写 len() 和 get()
• 从 CSV 加载：使用 pandas 加载节点和边表，构建连续索引映射，并组装成 Data 或 HeteroData
• 从 NetworkX 加载：from_networkx(G) 可直接转换 NetworkX 图
• 从 scipy 稀疏矩阵加载：from_scipy_sparse_matrix(adj) 提取 edge_index

阅读 references/custom_datasets.md，获取涵盖所有模式的完整示例、使用编码器加载 CSV 的示例，以及 MovieLens 演练。

可解释性

PyG 提供 torch_geometric.explain，用于解释 GNN 的预测：

from torch_geometric.explain import Explainer, GNNExplainer

explainer = Explainer(
    model=model,
    algorithm=GNNExplainer(epochs=200),
    explanation_type='model',
    node_mask_type='attributes',
    edge_mask_type='object',
    model_config=dict(
        mode='multiclass_classification',
        task_level='node',
        return_type='log_probs',
    ),
)

explanation = explainer(data.x, data.edge_index, index=10)
explanation.visualize_graph()           # Important subgraph
explanation.visualize_feature_importance(top_k=10)  # Feature importance

可用算法：GNNExplainer（基于优化）、PGExplainer（参数化、经过训练）、CaptumExplainer（通过 Captum 实现的基于梯度的方法）、AttentionExplainer（注意力权重）。适用于同构图和异构图。

阅读 references/explainability.md，了解所有算法、异构图解释、评估指标和 PGExplainer 训练。

常见陷阱

1. edge_index 形状：必须为 [2, num_edges]，而非 [num_edges, 2]。如有需要，请转置。
2. 忘记添加激活函数：卷积层不包含 ReLU 等激活函数，请手动添加。
3. 异构二部图中的自环：当源节点类型与目标节点类型不同时，不要使用 add_self_loops=True。请改用跳跃连接。
4. NeighborLoader 切片：只有前 batch.batch_size 个节点是种子节点。请相应地切片预测结果和标签。
5. 无向图：如果图是无向的，请在 edge_index 中包含两个方向的边，或使用 T.ToUndirected()。
6. 延迟初始化：输入通道为 -1 的模型需要先在 torch.no_grad() 下进行一次前向传播，以初始化参数，然后再开始训练。
7. 图任务的全局池化：使用 global_mean_pool(x, batch)（不要手动 reshape）将节点特征聚合到图级别。
8. num_neighbors 对齐：确保 len(num_neighbors) 等于 GNN 层数。跳数多于层数会浪费计算；跳数少于层数则会浪费模型容量。