torch-geometric

PyTorch Geometric (PyG)

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

安装

已针对 torch-geometric 2.7.x（2025 年 10 月）进行测试。需要 Python 3.10+ 和 PyTorch 2.6+。

# 1. Install PyTorch first (match your CUDA/CPU setup — see https://pytorch.org/get-started/locally/)
uv pip install torch

# 2. Core PyG (no extension wheels required for basic usage)
uv pip install torch_geometric

基本使用 PyG（从 PyG 2.3 起）不需要可选的加速算子（pyg-lib、torch-scatter、torch-sparse、torch-cluster）。检查 PyTorch 和 CUDA 版本后，从 PyG wheel index 安装匹配版本的 wheel：

python -c "import torch; print(torch.__version__, torch.version.cuda)"
# Then install wheels for your torch+CUDA combo, e.g.:
uv pip install pyg-lib torch-scatter torch-sparse torch-cluster \
  -f https://data.pyg.org/whl/torch-2.8.0+cu128.html

检查版本：

import torch_geometric
print(torch_geometric.__version__)

Conda： pyg conda channel 不再维护 PyTorch >2.5 的版本，请改用 uv pip install 和上面的 wheel index。

PyG 2.7 注意事项

PyG 2.7 不再支持 Python 3.9 和 PyTorch ≤2.5。PyTorch 2.6–2.8 的兼容性表请参阅 2.7.0 发布说明。torch_geometric.distributed 已弃用，请改用标准的 torch.distributed DDP（参见 references/scaling.md）。

核心概念

图数据：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() 汇总每个节点的消息（求和/求均值/求最大值）
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（源/邻居节点）来自动索引。因此，如果将 x=... 传递给 propagate，就可以在 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 — train(False) puts the model in inference mode (disables dropout/BN)
model.train(False)
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)（而不是手动重塑），将节点特征聚合为图级特征。
8. num_neighbors 对齐：确保 len(num_neighbors) 等于 GNN 层数。跳数多于层数会浪费计算；跳数较少则会浪费模型容量。

Agent 操作流程

1. 检查环境。 确认 Python 环境和库版本（python -c "import pkg; print(pkg.__version__)"），并检查数据的形状、类型和缺失值。
2. 明确输入。 根据数据或用户提供的信息确认格式、标识符和参数。任何会改变结果的值都应先询问，不要猜测。
3. 先运行小规模版本。 先在样本或单个折上运行，并检查运行时间和内存。
4. 使用上述说明和参考资料执行完整任务。
5. 验证结果。 使用留出数据、固定随机种子和适当的指标；检查是否存在数据泄漏；报告不确定性（置信区间、不同随机种子的标准差）。
6. 报告。 说明运行了什么（版本、命令、参数）、检查了什么，以及仍有哪些不确定之处。

如果出现这种情况	请执行以下操作
内存不足或运行非常缓慢	对数据进行子采样，使用分块或惰性计算，或缩小模型规模，并告知用户所做的更改。
这些说明中的函数、标志或端点在已安装版本中不存在	查阅已安装版本自己的文档（help()、--help、官方文档），进行适配，并告知用户。切勿虚构 API。
必需的输入、标识符或参数不明确	询问用户，或在运行前明确说明假设。

诚信规则

• 切勿捏造结果、参数、标识符、引文或统计数据。如果无法运行或验证某项内容，请如实说明。
• 切勿报告本次会话中未计算的指标；展示生成每个数值的代码路径。
• 请将此处与版本相关的细节视为可能已过时，并根据已安装版本的官方文档进行确认。
• 对于会产生费用、占用共享 GPU 或云配额、涉及个人或患者数据，或无法撤销的操作，须先征得同意。

相关技能

• networkx：使用 NetworkX 在 Python 中创建、分析和可视化复杂网络与图。
• torchdrug：构建和排查 TorchDrug 0.2.1 工作流，涵盖分子图、属性预测、自监督预训练、分子生成……