PyTorch Geometric (PyG)
PyG 是基于 PyTorch 构建的标准图神经网络库。它提供了图的数据结构、60 多种 GNN 层实现、可扩展的小批量训练以及对异构图的全面支持。
安装
已在 torch-geometric 2.7.x（2025年10月）上测试通过。需要 Python 3.10+ 和 PyTorch 2.6+。
# 1. 先安装 PyTorch（根据你的 CUDA/CPU 环境选择 — 请参见 https://pytorch.org/get-started/locally/）
uv pip install torch
# 2. PyG 核心库（基础使用无需扩展 wheel）
uv pip install torch_geometric
可选加速操作（pyg-lib、torch-scatter、torch-sparse、torch-cluster）并非 PyG 基础使用所必需（自 PyG 2.3 起）。在确认你的 PyTorch 和 CUDA 版本后，可从 PyG wheel 索引 安装版本匹配的 wheel：
python -c "import torch; print(torch.__version__, torch.version.cuda)"
# 然后为你的 torch+CUDA 组合安装 wheel，例如：
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 频道已不再维护 PyTorch >2.5 的版本 — 请改用 uv pip install 和上述 wheel 索引。
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 格式，dtype=torch.long
edge_attr=edge_features, # [num_edges, num_edge_features]
y=labels, # 节点级 [num_nodes, *] 或图级 [1, *]
pos=positions, # [num_nodes, num_dimensions]（用于点云/空间数据）
)
edge_index 格式至关重要：它是一个 [2, num_edges] 的张量，其中 edge_index[0] = 源节点，edge_index[1] = 目标节点。它不是元组列表。如果你有边对以行形式存储，请转置并调用 .contiguous()：
# 如果边是 [[src1, dst1], [src2, dst2], ...] — 先转置：
edge_index = edge_pairs.t().contiguous()
对于无向图，需包含两个方向：边 (0,1) 需要在 edge_index 中包含 [0,1] 和 [1,0]。
对于异构图，请使用 HeteroData — 请参见下面的异构图章节。
数据集
PyG 捆绑了许多标准数据集，可自动下载并预处理：
from torch_geometric.datasets import Planetoid, TUDataset
# 单图节点分类（Cora、Citeseer、Pubmed）
dataset = Planetoid(root='./data', name='Cora')
data = dataset[0] # 单个图，带有 train/val/test 掩码
# 多图分类（ENZYMES、MUTAG、IMDB-BINARY 等）
dataset = TUDataset(root='./data', name='ENZYMES')
# dataset[0], dataset[1], ... 是单个图
按任务划分的常见数据集：
- 节点分类：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
# 常用变换
T.NormalizeFeatures() # 按行归一化节点特征，使每行之和为 1
T.ToUndirected() # 添加反向边，使图成为无向图
T.AddSelfLoops() # 添加自循环边
T.KNNGraph(k=6) # 从点云位置构建 k-NN 图
T.RandomJitter(0.01) # 对位置进行随机噪声增强
T.Compose([...]) # 链式组合多个变换
# 作为 pre_transform（执行一次，保存到磁盘）或 transform（每次访问时执行）应用
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 | 点云、动态图 | 基于边特征的 MLP (x_i, x_j - x_i) |
RGCNConv | 多种关系类型的异构图 | 关系特定的权重矩阵 |
HGTConv | 异构图 | 类型特定的注意力 |
所有卷积层至少接受 (x, edge_index)。许多卷积层还接受 edge_attr 用于边特征。
延迟初始化
对于输入通道数使用 -1，让 PyG 自动推断维度——这在异构模型中特别有用：
conv = SAGEConv((-1, -1), 64) # 输入维度在首次前向传播时推断
# 初始化延迟模块：
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。框架如下：
propagate()协调消息传递message()定义沿每条边流动的信息（phi 函数）aggregate()在每个节点处聚合消息（sum/mean/max）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" 或 "max"
self.lin = torch.nn.Linear(in_channels, out_channels)
def forward(self, x, edge_index):
# 消息传递前的预处理
x = self.lin(x)
# 开始消息传递
return self.propagate(edge_index, x=x)
def message(self, x_j):
# x_j: 每条边源节点的特征 [num_edges, features]
# _j 后缀自动索引源节点，_i 索引目标节点
return x_j
_i / _j 约定：传递给 propagate() 的任何张量都可以通过在 message() 签名中添加 _i（目标/中心节点）或 _j（源/邻居节点）来自动索引。因此，如果你将 x=... 传递给 propagate，你可以在 message() 中访问 x_i 和 x_j。
完整的 GCN 和 EdgeConv 实现示例，请阅读 references/message_passing.md。
任务特定模式
节点分类
# 单图全批量训练（例如 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()
# 评估 — train(False) 将模型置于推理模式（禁用 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)
# 训练循环
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)
# 编码节点，然后计算边的分数
z = model.encode(train_data.x, train_data.edge_index)
# 正边
pos_score = (z[train_data.edge_label_index[0]] * z[train_data.edge_label_index[1]]).sum(dim=1)
完整的链接预测指南，包括 GAE/VGAE 自编码器、完整训练循环、适用于大图的 LinkNeighborLoader、异构链接预测和评估指标，请阅读 references/link_prediction.md。
大规模图扩展
对于无法放入 GPU 内存的图，可使用 NeighborLoader 进行邻域采样：
from torch_geometric.loader import NeighborLoader
train_loader = NeighborLoader(
data,
num_neighbors=[15, 10], # 在第 1 跳中采样 15 个邻居，在第 2 跳中采样 10 个
batch_size=128, # 每个批次中的种子节点数量
input_nodes=data.train_mask, # 要采样的节点
shuffle=True,
)
for batch in train_loader:
batch = batch.to(device)
out = model(batch.x, batch.edge_index)
# 只使用前 batch_size 个节点计算损失（这些是种子节点）
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()
# 节点特征 — 按节点类型字符串索引
data['user'].x = torch.randn(1000, 64)
data['movie'].x = torch.randn(500, 128)
# 边索引 — 按 (src_type, edge_type, dst_type) 三元组索引
data['user', 'rates', 'movie'].edge_index = torch.randint(0, 500, (2, 3000))
data['user', 'follows', 'user'].edge_index = torch.randint(0, 1000, (2, 5000))
# 便捷字典访问
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')
# 现在接受字典输入：
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, 64), 64),
('paper', 'revised_by', 'author'): SAGEConv((64, -1), 64),
}, aggr='sum')
# 然后在 forward 中使用
x_dict = conv(x_dict, edge_index_dict)
3. 手动实现 — 对于需要跨边类型精细控制的复杂模型，你可以在一个模块中为每个边类型编写自定义的 message/aggregation 逻辑。
详细教程包括异构 NeighborLoader、链接预测示例、元路径、HeteroData 的 to_homogeneous() 转换和高级 HGTLoader/HAN 实现，请阅读 references/heterogeneous.md。
常见陷阱
训练/评估模式与 Dropout
始终使用 model.train() 和 model.train(False)（或 model.eval()）正确控制 dropout 和批归一化。不要在推理时手动设置 model.training = False 或依赖 torch.no_grad()。
edge_label_index 与 edge_label
RandomLinkSplit 会创建以下属性：
edge_label_index：[2, num_edges]— 要打分的边（始终存在）edge_label：[num_edges]— 标签（1 表示正边，0 表示负边）
这与 model.encode() 中由 edge_index 定义的图拓扑是分开的。
设备管理
PyG 不会自动将数据移动到 GPU。请手动执行：
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
data = data.to(device)
model = model.to(device)
对于 DataLoader 迭代，将每个 batch 移到设备上：
for batch in loader:
batch = batch.to(device)
更多帮助，请运行 pyg -h 或访问 PyG 文档。
参考文献
该技能引用了以下文件以获得更深入的指导：
references/message_passing.md— 通过 MessagePassing 实现的自定义 GNN 层references/link_prediction.md— 完整的链接预测管道references/scaling.md— 大规模图训练（采样、集群、多 GPU、torch.compile）references/heterogeneous.md— 异构图模型、元路径、转换器