torch-geometric-graph-neural-networks

PyTorch Geometric（PyG）——图神经网络

概述

PyTorch Geometric 是一个基于 PyTorch 构建的库，用于开发和训练图神经网络（GNN）。它提供 40 多种卷积层、通过块对角邻接矩阵进行的小批量处理、适用于大规模图的邻居采样，以及对多类型节点/边网络的异构图支持。

使用场景

• 在引文网络、社交网络或生物网络上进行节点分类
• 图级分类（分子活性、蛋白质功能）
• 链接预测（知识图谱、推荐系统）
• 分子属性预测（药物发现、量子化学）
• 3D 点云处理和网格分析
• 使用邻居采样进行大规模图学习（>100K 个节点）
• 包含多种节点/边类型的异构图
• 非图深度学习 → 直接使用 PyTorch
• 传统图算法（最短路径、中心性） → 使用 NetworkX

先决条件

pip install torch torch_geometric
# Optional sparse operations (recommended):
# pip install pyg_lib torch_scatter torch_sparse torch_cluster

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

快速入门

from torch_geometric.datasets import Planetoid
from torch_geometric.nn import GCNConv
import torch, torch.nn.functional as F

dataset = Planetoid(root='/tmp/Cora', name='Cora')
data = dataset[0]

class GCN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = GCNConv(dataset.num_features, 16)
        self.conv2 = GCNConv(16, dataset.num_classes)
    def forward(self, data):
        x = F.relu(self.conv1(data.x, data.edge_index))
        return self.conv2(x, data.edge_index)

model = GCN()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
for epoch in range(200):
    model.train(); optimizer.zero_grad()
    F.cross_entropy(model(data)[data.train_mask], data.y[data.train_mask]).backward()
    optimizer.step()

model.eval()
pred = model(data).argmax(dim=1)
acc = (pred[data.test_mask] == data.y[data.test_mask]).float().mean()
print(f'Test Accuracy: {acc:.4f}')  # ~0.81

核心 API

1. 数据表示

import torch
from torch_geometric.data import Data

# Create a graph: 3 nodes, 4 edges (undirected)
edge_index = torch.tensor([[0, 1, 1, 2],
                           [1, 0, 2, 1]], dtype=torch.long)
x = torch.randn(3, 16)  # Node features [num_nodes, features]
y = torch.tensor([0, 1, 0])  # Node labels

data = Data(x=x, edge_index=edge_index, y=y)
print(f'Nodes: {data.num_nodes}, Edges: {data.num_edges}')
print(f'Features: {data.num_node_features}')
print(f'Has self-loops: {data.has_self_loops()}')
print(f'Is undirected: {data.is_undirected()}')

# Optional attributes
data.edge_attr = torch.randn(4, 8)   # Edge features [num_edges, features]
data.pos = torch.randn(3, 3)          # Node positions (3D)
data.train_mask = torch.tensor([True, True, False])  # Custom masks

# Mini-batch processing — graphs concatenated as block-diagonal
from torch_geometric.loader import DataLoader

loader = DataLoader(dataset, batch_size=32, shuffle=True)
for batch in loader:
    print(f'Graphs: {batch.num_graphs}, Nodes: {batch.num_nodes}')
    # batch.batch maps each node → its source graph index
    # No padding needed — computationally efficient

2. 卷积层

from torch_geometric.nn import GCNConv, GATConv, SAGEConv, GINConv
import torch.nn as nn

# GCNConv — spectral graph convolution (baseline)
conv = GCNConv(in_channels=16, out_channels=32)
# Supports: edge_weight, SparseTensor, Bipartite, Lazy init

# GATConv — attention-based neighbor weighting
conv = GATConv(16, 32, heads=8, dropout=0.6)
# Output: [N, heads * out_channels] (concat) or [N, out_channels] (concat=False)

# SAGEConv — inductive learning via sampling
conv = SAGEConv(16, 32, aggr='mean')  # 'mean', 'max', 'lstm'

# GINConv — maximally powerful for graph isomorphism
nn_module = nn.Sequential(nn.Linear(16, 32), nn.ReLU(), nn.Linear(32, 32))
conv = GINConv(nn_module)

# TransformerConv — graph transformer
from torch_geometric.nn import TransformerConv
conv = TransformerConv(16, 32, heads=8, beta=True)

# All layers: x_out = conv(x, edge_index)
x_out = conv(x, edge_index)
print(f'Output shape: {x_out.shape}')  # [num_nodes, out_channels]

3. 自定义消息传递

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

class CustomConv(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):
        edge_index, _ = add_self_loops(edge_index, num_nodes=x.size(0))
        x = self.lin(x)

        # Degree-based normalization
        row, col = edge_index
        deg = degree(col, x.size(0), dtype=x.dtype)
        norm = deg.pow(-0.5)
        norm = norm[row] * norm[col]

        return self.propagate(edge_index, x=x, norm=norm)

    def message(self, x_j, norm):
        # x_j: source node features (automatic via _j suffix)
        return norm.view(-1, 1) * x_j

# Key methods: forward(), message(), aggregate(), update()
# _i suffix → target node, _j suffix → source node

4. 池化与图级读出

from torch_geometric.nn import (
    global_mean_pool, global_max_pool, global_add_pool,
    TopKPooling, SAGPooling
)

# Global pooling: node features → graph-level representation
x_graph = global_mean_pool(x, batch)  # [num_graphs, features]

# Hierarchical pooling: coarsen graph
pool = TopKPooling(64, ratio=0.8)  # Keep top 80% nodes
x, edge_index, _, batch, _, _ = pool(x, edge_index, None, batch)

# Graph classification model
class GraphClassifier(torch.nn.Module):
    def __init__(self, num_features, num_classes):
        super().__init__()
        self.conv1 = GCNConv(num_features, 64)
        self.conv2 = GCNConv(64, 64)
        self.pool = TopKPooling(64, ratio=0.8)
        self.lin = torch.nn.Linear(64, num_classes)

    def forward(self, data):
        x, edge_index, batch = data.x, data.edge_index, data.batch
        x = F.relu(self.conv1(x, edge_index))
        x, edge_index, _, batch, _, _ = self.pool(x, edge_index, None, batch)
        x = F.relu(self.conv2(x, edge_index))
        x = global_mean_pool(x, batch)
        return self.lin(x)

5. 异构图

from torch_geometric.data import HeteroData
from torch_geometric.nn import HeteroConv, GCNConv, SAGEConv, to_hetero

# Create heterogeneous graph
data = HeteroData()
data['paper'].x = torch.randn(100, 128)
data['author'].x = torch.randn(200, 64)
data['author', 'writes', 'paper'].edge_index = torch.randint(0, 200, (2, 500))
data['paper', 'cites', 'paper'].edge_index = torch.randint(0, 100, (2, 300))
print(data)  # Shows all node/edge types

# Method 1: Auto-convert homogeneous model
model = GCN(...)
model = to_hetero(model, data.metadata(), aggr='sum')
out = model(data.x_dict, data.edge_index_dict)

# Method 2: Custom per-edge-type convolutions
class HeteroGNN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = HeteroConv({
            ('paper', 'cites', 'paper'): GCNConv(-1, 64),
            ('author', 'writes', 'paper'): SAGEConv((-1, -1), 64),
        }, aggr='sum')

    def forward(self, x_dict, edge_index_dict):
        x_dict = self.conv1(x_dict, edge_index_dict)
        return {k: F.relu(v) for k, v in x_dict.items()}

6. 变换与预处理

from torch_geometric.transforms import (
    NormalizeFeatures, AddSelfLoops, ToUndirected,
    RandomNodeSplit, RandomLinkSplit, Compose,
    KNNGraph, RadiusGraph, AddLaplacianEigenvectorPE
)

# Single transform
dataset = Planetoid(root='/tmp/Cora', name='Cora', transform=NormalizeFeatures())

# Compose multiple transforms
transform = Compose([
    ToUndirected(),
    AddSelfLoops(),
    NormalizeFeatures(),
])

# Data splitting
node_split = RandomNodeSplit(num_val=0.1, num_test=0.2)
link_split = RandomLinkSplit(num_val=0.1, num_test=0.2, is_undirected=True)

# Point cloud → graph
pc_transform = Compose([KNNGraph(k=6), NormalizeFeatures()])

# Positional encodings (for Graph Transformers)
pe_transform = AddLaplacianEigenvectorPE(k=10)

核心概念

层选择指南

任务	层	主要特性
基线 / 通用	GCNConv	谱方法、缓存、edge_weight
邻居重要性可变	GATConv / GATv2Conv	多头注意力
大规模归纳学习	SAGEConv	适合采样，支持 mean/max/lstm 聚合
图分类	GINConv	WL 测试中具有最大表达能力
长距离依赖	TransformerConv	图 Transformer
谱滤波	ChebConv	切比雪夫多项式，K 跳
丰富的边特征	NNConv	边神经网络处理 edge_attr
分子 / 3D 结构	SchNet、DimeNet	连续滤波器、角度
异构 / 多关系	RGCNConv、HGTConv	多种边类型
点云	EdgeConv、PointNetConv	动态图、局部特征
深层 GNN（避免过平滑）	APPNP + PairNorm	分离传播过程

数据流架构

• edge_index：[2, num_edges] COO 格式。第 0 行 = 源节点，第 1 行 = 目标节点
• 小批量：块对角邻接矩阵 + batch 向量，用于映射节点 → 图。不进行填充
• 邻居采样：NeighborLoader 为每个种子节点采样 K 跳子图。输出为有向图，且已重新标记节点
• 异构图：x_dict（每种类型的特征）、edge_index_dict（每种关系的边）、用于描述模式的 metadata()

聚合选项

聚合方式	类	使用场景
求和	SumAggregation	对计数敏感的任务
均值	MeanAggregation	不受节点度影响
最大值	MaxAggregation	显著特征检测
Softmax	SoftmaxAggregation(learn=True)	可学习的注意力
多种方式	MultiAggregation(['mean','max','std'])	组合信号

常见工作流

1. 节点分类（全图）

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

dataset = Planetoid(root='/tmp/Cora', name='Cora')
data = dataset[0]

class GCN(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = GCNConv(dataset.num_features, 16)
        self.conv2 = GCNConv(16, dataset.num_classes)
    def forward(self, data):
        x = F.dropout(F.relu(self.conv1(data.x, data.edge_index)), p=0.5, training=self.training)
        return self.conv2(x, data.edge_index)

model = GCN()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)

# Training
for epoch in range(200):
    model.train(); optimizer.zero_grad()
    out = model(data)
    F.cross_entropy(out[data.train_mask], data.y[data.train_mask]).backward()
    optimizer.step()

# Evaluation
model.eval()
pred = model(data).argmax(dim=1)
acc = (pred[data.test_mask] == data.y[data.test_mask]).float().mean()
print(f'Test Accuracy: {acc:.4f}')

2. 图分类（小批量）

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

dataset = TUDataset(root='/tmp/ENZYMES', name='ENZYMES')
train_dataset = dataset[:int(0.8 * len(dataset))]
test_dataset = dataset[int(0.8 * len(dataset)):]
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)

class GraphNet(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = GCNConv(dataset.num_features, 64)
        self.conv2 = GCNConv(64, 64)
        self.lin = torch.nn.Linear(64, dataset.num_classes)
    def forward(self, data):
        x = F.relu(self.conv1(data.x, data.edge_index))
        x = F.relu(self.conv2(x, data.edge_index))
        x = global_mean_pool(x, data.batch)
        return self.lin(x)

model = GraphNet()
optimizer = torch.optim.Adam(model.parameters(), lr=0.01)
for epoch in range(100):
    model.train()
    for batch in train_loader:
        optimizer.zero_grad()
        F.cross_entropy(model(batch), batch.y).backward()
        optimizer.step()

3. 使用邻居采样进行大规模学习

from torch_geometric.loader import NeighborLoader

# Sample 25 1-hop and 10 2-hop neighbors per seed node
train_loader = NeighborLoader(
    data,
    num_neighbors=[25, 10],
    batch_size=128,
    input_nodes=data.train_mask,
)

model.train()
for batch in train_loader:
    optimizer.zero_grad()
    out = model(batch)
    # Only compute loss on seed nodes (first batch_size nodes)
    loss = F.cross_entropy(out[:batch.batch_size], batch.y[:batch.batch_size])
    loss.backward()
    optimizer.step()
# Note: output subgraphs are directed, indices relabeled 0..N-1

关键参数

参数	模块	默认值	范围	作用
in_channels	所有 Conv 层	—	int	输入特征维度
out_channels	所有 Conv 层	—	int	输出特征维度
heads	GATConv	1	1-16	注意力头数量
dropout	GATConv	0.0	0-0.8	注意力权重 dropout
aggr	MessagePassing	'add'	add/mean/max	邻居聚合方式
K	ChebConv	—	2-5	切比雪夫多项式阶数
num_neighbors	NeighborLoader	—	list[int]	每跳采样的邻居数量（例如，[25,10]）
batch_size	DataLoader	—	16-512	每批次的图或种子节点数量
ratio	TopKPooling	0.5	0.1-0.9	保留的节点比例
lr	Adam	—	1e-4 to 0.01	学习率
weight_decay	Adam	0	0 to 5e-3	L2 正则化

最佳实践

1. 从 GCNConv 开始：先使用 2 层 GCN 作为基线，再尝试复杂架构
2. 使用延迟初始化：将 -1 作为 in_channels 传入，以自动推断维度：GCNConv(-1, 64)
3. 归一化特征：对引文网络或社交网络应用 NormalizeFeatures() 转换
4. 反模式 — 层数过多：GNN 通常只需要 2-3 层。层数更深会导致过平滑。如果需要更深的网络，请使用 JumpingKnowledge 或 PairNorm
5. 传输到 GPU：将模型和数据都移到 GPU：model.to(device)、data.to(device)
6. 反模式 — 忽略 batch 向量：在图分类中，始终使用 global_mean_pool(x, batch) — 忘记传入 batch 会对所有图进行池化

常见示例

示例：模型可解释性（GNNExplainer）

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)
print(f'Important edges: {explanation.edge_mask.topk(5).indices}')
print(f'Important features: {explanation.node_mask[10].topk(5).indices}')

示例：自定义 InMemoryDataset

from torch_geometric.data import InMemoryDataset, Data

class MyDataset(InMemoryDataset):
    def __init__(self, root, transform=None, pre_transform=None):
        super().__init__(root, transform, pre_transform)
        self.load(self.processed_paths[0])

    @property
    def raw_file_names(self):
        return ['data.csv']

    @property
    def processed_file_names(self):
        return ['data.pt']

    def process(self):
        data_list = []
        # Build Data objects from raw files
        edge_index = torch.tensor([[0, 1], [1, 0]], dtype=torch.long)
        x = torch.randn(2, 16)
        data_list.append(Data(x=x, edge_index=edge_index, y=torch.tensor([0])))

        if self.pre_filter is not None:
            data_list = [d for d in data_list if self.pre_filter(d)]
        if self.pre_transform is not None:
            data_list = [self.pre_transform(d) for d in data_list]
        self.save(data_list, self.processed_paths[0])

示例：使用 JumpingKnowledge 的深层 GNN

from torch_geometric.nn import GCNConv, JumpingKnowledge, LayerNorm

class DeepGNN(torch.nn.Module):
    def __init__(self, in_ch, hidden, num_layers, out_ch):
        super().__init__()
        self.convs = torch.nn.ModuleList()
        self.norms = torch.nn.ModuleList()
        self.convs.append(GCNConv(in_ch, hidden))
        self.norms.append(LayerNorm(hidden))
        for _ in range(num_layers - 2):
            self.convs.append(GCNConv(hidden, hidden))
            self.norms.append(LayerNorm(hidden))
        self.convs.append(GCNConv(hidden, hidden))
        self.jk = JumpingKnowledge(mode='cat')
        self.lin = torch.nn.Linear(hidden * num_layers, out_ch)

    def forward(self, x, edge_index, batch):
        xs = []
        for conv, norm in zip(self.convs[:-1], self.norms):
            x = F.relu(norm(conv(x, edge_index)))
            xs.append(x)
        xs.append(self.convs[-1](x, edge_index))
        return self.lin(global_mean_pool(self.jk(xs), batch))

故障排除

问题	原因	解决方案
edge_index 形状错误	格式错误（应为 [2, E]）	确保使用 COO 格式：torch.tensor([[src...],[dst...]], dtype=torch.long)
大图导致 OOM	全图前向传播	使用 NeighborLoader 进行小批量训练
准确率低	过平滑（层数过多）	减少到 2-3 层，并添加 JumpingKnowledge 或 PairNorm
训练中出现 NaN	梯度爆炸	添加梯度裁剪、降低学习率，并检查特征缩放
图级输出错误	池化时缺少 batch	将 batch 张量传给 global_mean_pool(x, batch)
异构类型错误	节点/边类型不匹配	检查 data.metadata() 是否与模型定义相符
DataLoader 速度慢	图过大且未进行采样	使用 NeighborLoader 并设置合理的 num_neighbors（例如，[25,10]）
x 维度不匹配	多头注意力输出	对于 GATConv：输出为 heads*out_channels，除非设置 concat=False
稀疏操作导入错误	缺少可选依赖项	从 PyG wheels 安装 torch_scatter、torch_sparse
未应用预转换	数据集已处理	删除 processed/ 目录并重新加载

内置资源

• references/layers_transforms_reference.md — 40 多种卷积层的完整目录（GCN、GAT、SAGE、GIN、分子层、超图）、聚合算子、池化（全局 + 分层）、归一化层、预构建模型、自动编码器、知识图谱嵌入和实用层。变换目录：结构、特征、空间、数据增强、网格、专用。整合自原始 layers_reference.md（486 行）和 transforms_reference.md（680 行）。脚本功能（benchmark_model.py、create_gnn_template.py、visualize_graph.py）由 Core API 代码块和常见示例介绍
• references/datasets_catalog.md — 按领域整理的综合数据集目录：引文网络（Planetoid、Coauthor、Amazon）、图分类（TUDataset，120 多个基准）、分子数据集（QM9、ZINC、MoleculeNet）、社交数据集（Reddit、Twitch）、知识图谱（WordNet、FB15k）、异构数据集（OGB_MAG、MovieLens、DBLP）、时序数据集（JODIE）、3D 网格（ShapeNet、ModelNet）、OGB 集成。整合自原始 datasets_reference.md（575 行）

相关技能

• matplotlib-scientific-plotting — 可视化图结构、训练曲线和注意力权重

参考资料

• PyTorch Geometric 文档：https://pytorch-geometric.readthedocs.io/
• PyG GitHub：https://github.com/pyg-team/pytorch_geometric
• Fey 和 Lenssen（2019），《Fast Graph Representation Learning with PyTorch Geometric》，ICLR Workshop