- Authors
- Name
- Youngju Kim
- @fjvbn20031
- 图神经网络(GNN)完全指南
图神经网络(GNN)完全指南
社交网络、分子结构、知识图谱、推荐系统 — 现实世界中无数的数据都以图的形式呈现。图神经网络(Graph Neural Network, GNN)正是用深度学习处理这类非欧几里得(non-Euclidean)数据的核心工具。本指南将系统梳理从图论基础到最新 GNN 架构、再到基于 PyTorch Geometric 的实战实现的完整内容。
1. 图论基础
图的定义
图 G 由节点(Node)集合 V 与边(Edge)集合 E 构成，记作 G = (V, E)。节点表示实体(entity)，边表示实体之间的关系。
- 节点(Vertex/Node)：表示实体。例：用户、原子、论文
- 边(Edge)：表示关系。例：好友关系、化学键、引用关系
- 节点特征(Node Feature)：连接到每个节点的特征向量
- 边特征(Edge Feature)：连接到每条边的特征向量
有向图/无向图
import networkx as nx
import matplotlib.pyplot as plt
import numpy as np
# 无向图 (Undirected Graph)
G_undirected = nx.Graph()
G_undirected.add_edges_from([(0, 1), (1, 2), (2, 3), (3, 0), (0, 2)])
# 有向图 (Directed Graph)
G_directed = nx.DiGraph()
G_directed.add_edges_from([(0, 1), (1, 2), (2, 0), (0, 3)])
print(f"无向图 - 节点数: {G_undirected.number_of_nodes()}, 边数: {G_undirected.number_of_edges()}")
print(f"有向图 - 节点数: {G_directed.number_of_nodes()}, 边数: {G_directed.number_of_edges()}")
邻接矩阵与边列表
import torch
import numpy as np
# 邻接矩阵 (Adjacency Matrix)
# A[i][j] = 1 表示节点 i 与 j 之间存在边
adj_matrix = torch.tensor([
[0, 1, 1, 0],
[1, 0, 1, 0],
[1, 1, 0, 1],
[0, 0, 1, 0]
], dtype=torch.float32)
# 边列表 (Edge Index) - PyG 使用的格式
# shape: (2, num_edges) - 第一行: 源节点, 第二行: 目标节点
edge_index = torch.tensor([
[0, 0, 1, 1, 2, 2, 2, 3], # 源节点
[1, 2, 0, 2, 0, 1, 3, 2] # 目标节点
], dtype=torch.long)
print(f"邻接矩阵大小: {adj_matrix.shape}") # (4, 4)
print(f"边列表大小: {edge_index.shape}") # (2, 8)
# 邻接矩阵 -> 边列表 转换
def adj_to_edge_index(adj):
"""将邻接矩阵转换为边索引"""
row, col = torch.where(adj > 0)
return torch.stack([row, col], dim=0)
converted = adj_to_edge_index(adj_matrix)
print(f"转换后的边列表:\n{converted}")
图的特性
import networkx as nx
import numpy as np
def analyze_graph(G):
"""分析图的主要特性"""
# 度数 (Degree)
degrees = dict(G.degree())
avg_degree = np.mean(list(degrees.values()))
# 聚类系数 (Clustering Coefficient)
clustering = nx.average_clustering(G)
# 平均路径长度 (Average Path Length)
if nx.is_connected(G):
avg_path = nx.average_shortest_path_length(G)
else:
# 使用最大的连通分量
largest_cc = max(nx.connected_components(G), key=len)
subgraph = G.subgraph(largest_cc)
avg_path = nx.average_shortest_path_length(subgraph)
# 中心性 (Centrality)
betweenness = nx.betweenness_centrality(G)
pagerank = nx.pagerank(G)
print(f"节点数: {G.number_of_nodes()}")
print(f"边数: {G.number_of_edges()}")
print(f"平均度数: {avg_degree:.2f}")
print(f"聚类系数: {clustering:.3f}")
print(f"平均路径长度: {avg_path:.2f}")
return {
"degrees": degrees,
"clustering": clustering,
"avg_path": avg_path,
"betweenness": betweenness,
"pagerank": pagerank
}
# 社交网络示例 (Karate Club)
G = nx.karate_club_graph()
stats = analyze_graph(G)
现实世界中的图
| 领域 | 节点 | 边 | 任务 |
|---|---|---|---|
| 社交网络 | 用户 | 好友关系 | 社区发现 |
| 分子结构 | 原子 | 化学键 | 分子性质预测 |
| 知识图谱 | 实体 | 关系 | 链接预测 |
| 引用网络 | 论文 | 引用关系 | 节点分类 |
| 交通网络 | 路口 | 道路 | 路径预测 |
| 推荐系统 | 用户/物品 | 交互 | 推荐 |
2. 图机器学习的动机
为什么 CNN/RNN 力不从心？
传统的 CNN 以网格(grid)结构为前提。图像的像素规则地排列在 2D 网格上，因此卷积运算能够自然地发挥作用。RNN 则假设数据具有序列(sequence)结构。
但图具有以下特性：
- 非规则结构：每个节点的邻居数各不相同
- 没有固定顺序：节点具有排列不变性(Permutation Invariance)
- 全局依赖性：距离较远的节点之间也可能相互影响
# 图数据的特性说明
# 图像: 固定大小的网格
image = torch.randn(3, 224, 224) # 通道, 高度, 宽度
# 序列: 有顺序的数据
sequence = torch.randn(100, 512) # 序列长度, 特征维度
# 图: 可变的邻居结构
# 节点特征: (num_nodes, feature_dim)
node_features = torch.randn(34, 16) # 34个节点, 16维特征
# 边: (2, num_edges) - 稀疏连接
edge_index = torch.randint(0, 34, (2, 78))
消息传递范式
所有 GNN 的基本原理都是消息传递(Message Passing)。每个节点从邻居节点接收消息，并以此更新自己的表示。
消息传递神经网络(MPNN)框架：
- 消息计算：针对边 (u, v)，计算从节点 u 传递给 v 的消息
- 聚合：每个节点将所有邻居消息合并
- 更新：用聚合后的消息更新节点表示
m_v^(l) = AGGREGATE({h_u^(l-1) : u in N(v)})
h_v^(l) = UPDATE(h_v^(l-1), m_v^(l))
其中 N(v) 是节点 v 的邻居集合。
3. GNN 基础公式
聚合(Aggregation)与更新(Update)
import torch
import torch.nn as nn
from torch_scatter import scatter_mean, scatter_sum, scatter_max
def manual_message_passing(node_features, edge_index, aggregation="mean"):
"""
手动实现的消息传递
node_features: (N, F) - N个节点, F维特征
edge_index: (2, E) - E条边
"""
src, dst = edge_index[0], edge_index[1]
num_nodes = node_features.size(0)
# 使用源节点的特征作为消息
messages = node_features[src] # (E, F)
if aggregation == "mean":
# 按目标节点求平均
aggregated = scatter_mean(messages, dst, dim=0, dim_size=num_nodes)
elif aggregation == "sum":
aggregated = scatter_sum(messages, dst, dim=0, dim_size=num_nodes)
elif aggregation == "max":
aggregated, _ = scatter_max(messages, dst, dim=0, dim_size=num_nodes)
# 更新: 原始特征 + 聚合后的消息
updated = node_features + aggregated
return updated
# 示例
N, F = 6, 8
node_features = torch.randn(N, F)
edge_index = torch.tensor([[0,1,2,3,4,0,1], [1,2,3,4,0,3,4]])
output = manual_message_passing(node_features, edge_index, "mean")
print(f"Input shape: {node_features.shape}")
print(f"Output shape: {output.shape}")
4. 主要 GNN 架构
GCN (Graph Convolutional Network)
Kipf & Welling (2017) 提出的 GCN 从谱图理论(spectral graph theory)出发，推导出了高效的逐层传播规则。
逐层传播规则：
使用归一化邻接矩阵的传播: tilde A = D^(-1/2) _ (A + I) _ D^(-1/2)
H^(l+1) = sigma(tilde A _ H^(l) _ W^(l))
其中 D 是度矩阵(degree matrix)，I 是单位矩阵，W 是可学习的权重。
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import GCNConv
from torch_geometric.data import Data
from torch_geometric.datasets import Planetoid
from torch_geometric.transforms import NormalizeFeatures
# 加载数据集
dataset = Planetoid(root='/tmp/Cora', name='Cora', transform=NormalizeFeatures())
data = dataset[0]
print(f"节点数: {data.num_nodes}")
print(f"边数: {data.num_edges}")
print(f"节点特征维度: {data.num_node_features}")
print(f"类别数: {dataset.num_classes}")
print(f"训练节点数: {data.train_mask.sum().item()}")
print(f"验证节点数: {data.val_mask.sum().item()}")
print(f"测试节点数: {data.test_mask.sum().item()}")
class GCN(nn.Module):
"""Graph Convolutional Network"""
def __init__(self, in_channels, hidden_channels, out_channels, dropout=0.5):
super().__init__()
self.conv1 = GCNConv(in_channels, hidden_channels)
self.conv2 = GCNConv(hidden_channels, out_channels)
self.dropout = dropout
def forward(self, x, edge_index):
# 第一个 GCN 层
x = self.conv1(x, edge_index)
x = F.relu(x)
x = F.dropout(x, p=self.dropout, training=self.training)
# 第二个 GCN 层
x = self.conv2(x, edge_index)
return x
# 设置模型与优化器
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = GCN(
in_channels=dataset.num_features,
hidden_channels=64,
out_channels=dataset.num_classes
).to(device)
data = data.to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
def train_gcn():
model.train()
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()
return loss.item()
def test_gcn():
model.eval()
with torch.no_grad():
out = model(data.x, data.edge_index)
pred = out.argmax(dim=1)
results = {}
for split, mask in [("train", data.train_mask),
("val", data.val_mask),
("test", data.test_mask)]:
correct = pred[mask].eq(data.y[mask]).sum().item()
results[split] = correct / mask.sum().item()
return results
# 训练循环
best_val_acc = 0
for epoch in range(200):
loss = train_gcn()
accs = test_gcn()
if accs["val"] > best_val_acc:
best_val_acc = accs["val"]
if (epoch + 1) % 50 == 0:
print(f"Epoch {epoch+1:03d} | Loss: {loss:.4f} | "
f"Train: {accs['train']:.4f} | Val: {accs['val']:.4f} | "
f"Test: {accs['test']:.4f}")
GCN 手动实现
import torch
import torch.nn as nn
import torch.nn.functional as F
class ManualGCNLayer(nn.Module):
"""GCN 层的手动实现 - 用于理解内部运作原理"""
def __init__(self, in_features, out_features):
super().__init__()
self.weight = nn.Parameter(torch.FloatTensor(in_features, out_features))
self.bias = nn.Parameter(torch.FloatTensor(out_features))
self.reset_parameters()
def reset_parameters(self):
nn.init.xavier_uniform_(self.weight)
nn.init.zeros_(self.bias)
def forward(self, x, adj):
"""
x: 节点特征 (N, F_in)
adj: 归一化后的邻接矩阵 (N, N)
"""
# 线性变换: X * W
support = x @ self.weight
# 图卷积: A_hat * X * W
output = adj @ support + self.bias
return output
@staticmethod
def normalize_adjacency(adj):
"""D^(-1/2) * A * D^(-1/2) 归一化"""
# 添加自环
N = adj.size(0)
adj_hat = adj + torch.eye(N, device=adj.device)
# 计算度矩阵
deg = adj_hat.sum(dim=1)
d_inv_sqrt = torch.diag(deg.pow(-0.5))
# 归一化
adj_normalized = d_inv_sqrt @ adj_hat @ d_inv_sqrt
return adj_normalized
GraphSAGE (归纳式学习)
GraphSAGE 专为归纳式(inductive)学习而设计。它不使用整张图，而是通过采样邻居实现小批量(mini-batch)训练。
from torch_geometric.nn import SAGEConv
import torch
import torch.nn as nn
import torch.nn.functional as F
class GraphSAGE(nn.Module):
"""GraphSAGE - 归纳式表示学习"""
def __init__(self, in_channels, hidden_channels, out_channels,
num_layers=3, dropout=0.5, aggr="mean"):
super().__init__()
self.dropout = dropout
self.convs = nn.ModuleList()
self.convs.append(SAGEConv(in_channels, hidden_channels, aggr=aggr))
for _ in range(num_layers - 2):
self.convs.append(SAGEConv(hidden_channels, hidden_channels, aggr=aggr))
self.convs.append(SAGEConv(hidden_channels, out_channels, aggr=aggr))
self.bns = nn.ModuleList([
nn.BatchNorm1d(hidden_channels)
for _ in range(num_layers - 1)
])
def forward(self, x, edge_index):
for i, conv in enumerate(self.convs[:-1]):
x = conv(x, edge_index)
x = self.bns[i](x)
x = F.relu(x)
x = F.dropout(x, p=self.dropout, training=self.training)
x = self.convs[-1](x, edge_index)
return x
# 使用邻居采样进行小批量训练
from torch_geometric.loader import NeighborLoader
# NeighborLoader: 每一层采样 num_neighbors 个邻居
train_loader = NeighborLoader(
data,
num_neighbors=[25, 10], # 2-hop: 第一跳采样 25 个, 第二跳采样 10 个
batch_size=256,
input_nodes=data.train_mask,
shuffle=True
)
model_sage = GraphSAGE(
in_channels=dataset.num_features,
hidden_channels=64,
out_channels=dataset.num_classes
).to(device)
optimizer_sage = torch.optim.Adam(model_sage.parameters(), lr=0.001)
def train_sage():
model_sage.train()
total_loss = 0
for batch in train_loader:
batch = batch.to(device)
optimizer_sage.zero_grad()
out = model_sage(batch.x, batch.edge_index)
# 只有批次最前面的 batch_size 个节点是训练节点
loss = F.cross_entropy(out[:batch.batch_size], batch.y[:batch.batch_size])
loss.backward()
optimizer_sage.step()
total_loss += loss.item()
return total_loss / len(train_loader)
GAT (Graph Attention Network)
GAT 使用注意力机制，为每个邻居赋予不同的权重，体现了「并非所有邻居都同等重要」这一直觉。
注意力系数计算：
注意力分数: e_ij = LeakyReLU(a^T [Wh_i || Wh_j])
Softmax 归一化: alpha_ij = exp(e_ij) / sum_k(exp(e_ik))
更新: h_i' = sigma(sum_j alpha_ij _ W _ h_j)
from torch_geometric.nn import GATConv, GATv2Conv
import torch
import torch.nn as nn
import torch.nn.functional as F
class GAT(nn.Module):
"""Graph Attention Network"""
def __init__(self, in_channels, hidden_channels, out_channels,
heads=8, dropout=0.6):
super().__init__()
self.dropout = dropout
# 第一层: 多头注意力
self.conv1 = GATConv(
in_channels,
hidden_channels,
heads=heads,
dropout=dropout,
concat=True # 拼接(concatenate)多个头
)
# 第二层: 对多个头取平均
self.conv2 = GATConv(
hidden_channels * heads,
out_channels,
heads=1,
dropout=dropout,
concat=False # 对头取平均
)
def forward(self, x, edge_index):
x = F.dropout(x, p=self.dropout, training=self.training)
x = self.conv1(x, edge_index)
x = F.elu(x)
x = F.dropout(x, p=self.dropout, training=self.training)
x = self.conv2(x, edge_index)
return x
class GATv2(nn.Module):
"""
GATv2 - 改进的注意力机制
GATv2 计算动态注意力, 表达能力更强
"""
def __init__(self, in_channels, hidden_channels, out_channels,
heads=8, dropout=0.6):
super().__init__()
self.conv1 = GATv2Conv(
in_channels,
hidden_channels,
heads=heads,
dropout=dropout,
concat=True
)
self.conv2 = GATv2Conv(
hidden_channels * heads,
out_channels,
heads=1,
dropout=dropout,
concat=False
)
self.dropout = dropout
def forward(self, x, edge_index):
x = F.dropout(x, p=self.dropout, training=self.training)
x = F.elu(self.conv1(x, edge_index))
x = F.dropout(x, p=self.dropout, training=self.training)
return self.conv2(x, edge_index)
# 训练 GAT
model_gat = GAT(
in_channels=dataset.num_features,
hidden_channels=8,
out_channels=dataset.num_classes,
heads=8
).to(device)
optimizer_gat = torch.optim.Adam(model_gat.parameters(), lr=0.005, weight_decay=5e-4)
Graph Transformer
Graph Transformer 将 Transformer 的全局注意力机制应用到图上。
from torch_geometric.nn import TransformerConv
import torch
import torch.nn as nn
import torch.nn.functional as F
class GraphTransformer(nn.Module):
"""Graph Transformer Layer"""
def __init__(self, in_channels, hidden_channels, out_channels,
heads=4, num_layers=3, dropout=0.3):
super().__init__()
self.dropout = dropout
self.convs = nn.ModuleList()
self.convs.append(
TransformerConv(in_channels, hidden_channels // heads, heads=heads,
dropout=dropout, beta=True)
)
for _ in range(num_layers - 2):
self.convs.append(
TransformerConv(hidden_channels, hidden_channels // heads,
heads=heads, dropout=dropout, beta=True)
)
self.convs.append(
TransformerConv(hidden_channels, out_channels // heads,
heads=heads, dropout=dropout, beta=True)
)
self.norms = nn.ModuleList([
nn.LayerNorm(hidden_channels) for _ in range(num_layers - 1)
])
def forward(self, x, edge_index):
for i, conv in enumerate(self.convs[:-1]):
x = conv(x, edge_index)
x = self.norms[i](x)
x = F.relu(x)
x = F.dropout(x, p=self.dropout, training=self.training)
return self.convs[-1](x, edge_index)
5. 图级别预测
如果说节点分类是对单个节点的预测，那么图分类就是对整张图的预测。例如：预测某个分子是否有毒。
Global Pooling
from torch_geometric.nn import (
global_mean_pool,
global_max_pool,
global_add_pool
)
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import GCNConv
class GraphClassifier(nn.Module):
"""图分类模型"""
def __init__(self, in_channels, hidden_channels, out_channels,
num_layers=3, dropout=0.5, pooling="mean"):
super().__init__()
self.dropout = dropout
self.pooling = pooling
self.convs = nn.ModuleList()
self.convs.append(GCNConv(in_channels, hidden_channels))
for _ in range(num_layers - 1):
self.convs.append(GCNConv(hidden_channels, hidden_channels))
self.bns = nn.ModuleList([
nn.BatchNorm1d(hidden_channels) for _ in range(num_layers)
])
# 图级别分类器
self.classifier = nn.Sequential(
nn.Linear(hidden_channels, hidden_channels),
nn.ReLU(),
nn.Dropout(dropout),
nn.Linear(hidden_channels, out_channels)
)
def forward(self, x, edge_index, batch):
"""
batch: 表示每个节点属于哪张图的索引向量
"""
# 节点嵌入
for conv, bn in zip(self.convs, self.bns):
x = conv(x, edge_index)
x = bn(x)
x = F.relu(x)
x = F.dropout(x, p=self.dropout, training=self.training)
# 图级别池化
if self.pooling == "mean":
x = global_mean_pool(x, batch)
elif self.pooling == "max":
x = global_max_pool(x, batch)
elif self.pooling == "sum":
x = global_add_pool(x, batch)
# 分类
return self.classifier(x)
DiffPool (Differentiable Pooling)
from torch_geometric.nn import dense_diff_pool
import torch
import torch.nn as nn
import torch.nn.functional as F
class DiffPoolLayer(nn.Module):
"""层次化图池化"""
def __init__(self, in_channels, hidden_channels, num_clusters):
super().__init__()
# GNN for node embedding
self.gnn_embed = nn.Sequential(
nn.Linear(in_channels, hidden_channels),
nn.ReLU()
)
# GNN for cluster assignment
self.gnn_pool = nn.Sequential(
nn.Linear(in_channels, num_clusters),
)
def forward(self, x, adj, mask=None):
embed = self.gnn_embed(x)
# Cluster assignment matrix
s = torch.softmax(self.gnn_pool(x), dim=-1)
# DiffPool
out, out_adj, link_loss, entropy_loss = dense_diff_pool(embed, adj, s, mask)
return out, out_adj, link_loss, entropy_loss
6. 链接预测
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import GCNConv
from torch_geometric.utils import negative_sampling
from torch_geometric.transforms import RandomLinkSplit
class LinkPredictor(nn.Module):
"""链接预测模型"""
def __init__(self, in_channels, hidden_channels, out_channels):
super().__init__()
# 节点嵌入编码器
self.encoder = nn.ModuleList([
GCNConv(in_channels, hidden_channels),
GCNConv(hidden_channels, out_channels)
])
# 边解码器
self.decoder = nn.Sequential(
nn.Linear(out_channels * 2, out_channels),
nn.ReLU(),
nn.Linear(out_channels, 1)
)
def encode(self, x, edge_index):
for i, conv in enumerate(self.encoder):
x = conv(x, edge_index)
if i < len(self.encoder) - 1:
x = F.relu(x)
return x
def decode(self, z, edge_index):
# 拼接源/目标节点嵌入
src, dst = edge_index
edge_feat = torch.cat([z[src], z[dst]], dim=1)
return self.decoder(edge_feat).squeeze()
def forward(self, x, edge_index, pos_edge_index, neg_edge_index):
z = self.encode(x, edge_index)
pos_pred = self.decode(z, pos_edge_index)
neg_pred = self.decode(z, neg_edge_index)
return pos_pred, neg_pred
def train_link_prediction(model, data, optimizer):
model.train()
optimizer.zero_grad()
# 节点嵌入
z = model.encode(data.x, data.edge_index)
# 正样本边
pos_edge = data.train_pos_edge_index
# 负样本边采样
neg_edge = negative_sampling(
edge_index=pos_edge,
num_nodes=data.num_nodes,
num_neg_samples=pos_edge.size(1)
)
pos_pred = model.decode(z, pos_edge)
neg_pred = model.decode(z, neg_edge)
# Binary cross-entropy loss
pred = torch.cat([pos_pred, neg_pred])
labels = torch.cat([
torch.ones(pos_pred.size(0)),
torch.zeros(neg_pred.size(0))
]).to(pred.device)
loss = F.binary_cross_entropy_with_logits(pred, labels)
loss.backward()
optimizer.step()
return loss.item()
7. PyTorch Geometric (PyG) 完全指南
安装
pip install torch-geometric
pip install torch-scatter torch-sparse -f https://data.pyg.org/whl/torch-2.1.0+cu121.html
Data 对象
from torch_geometric.data import Data
import torch
# 创建图数据
x = torch.randn(6, 3) # 6个节点, 3维特征
edge_index = torch.tensor([
[0, 1, 2, 3, 4, 0],
[1, 2, 3, 4, 0, 3]
], dtype=torch.long)
y = torch.tensor([0, 1, 0, 1, 0, 1]) # 节点标签
edge_attr = torch.randn(6, 2) # 边特征
data = Data(
x=x,
edge_index=edge_index,
y=y,
edge_attr=edge_attr
)
print(data)
print(f"节点数: {data.num_nodes}")
print(f"边数: {data.num_edges}")
print(f"节点特征维度: {data.num_node_features}")
print(f"边特征维度: {data.num_edge_features}")
print(f"是否存在 self-loop: {data.has_self_loops()}")
print(f"是否为有向图: {data.is_directed()}")
# 有效性检查
print(f"数据是否有效: {data.validate()}")
DataLoader 与小批量
from torch_geometric.data import Data, DataLoader
import torch
# 创建图数据集
dataset = []
for _ in range(100):
n = torch.randint(5, 20, (1,)).item() # 5~20个节点
e = torch.randint(10, 40, (1,)).item() # 10~40条边
data = Data(
x=torch.randn(n, 8),
edge_index=torch.randint(0, n, (2, e)),
y=torch.randint(0, 3, (1,)) # 图标签
)
dataset.append(data)
# DataLoader: 将多张图打包成一张不连续的大图批次
loader = DataLoader(dataset, batch_size=32, shuffle=True)
for batch in loader:
print(f"批次中的图数: {batch.num_graphs}")
print(f"节点总数: {batch.num_nodes}")
print(f"边总数: {batch.num_edges}")
print(f"batch 向量: {batch.batch.shape}") # 每个节点所属的图索引
break
内置数据集
from torch_geometric.datasets import (
Planetoid, # Cora, Citeseer, PubMed
TUDataset, # 分子数据集 (MUTAG, ENZYMES 等)
OGB, # Open Graph Benchmark
)
from torch_geometric.transforms import NormalizeFeatures, RandomNodeSplit
# Cora 引用网络
cora = Planetoid(root='/tmp/Cora', name='Cora', transform=NormalizeFeatures())
print(f"Cora - 节点数: {cora[0].num_nodes}, 边数: {cora[0].num_edges}")
# MUTAG 分子数据集
mutag = TUDataset(root='/tmp/TUDataset', name='MUTAG')
print(f"MUTAG - 图数: {len(mutag)}, 类别数: {mutag.num_classes}")
# Open Graph Benchmark (大规模)
try:
from ogb.nodeproppred import PygNodePropPredDataset
dataset_ogb = PygNodePropPredDataset(name='ogbn-arxiv')
split_idx = dataset_ogb.get_idx_split()
data_ogb = dataset_ogb[0]
print(f"OGB-Arxiv - 节点数: {data_ogb.num_nodes}, 边数: {data_ogb.num_edges}")
except ImportError:
print("未安装 ogb 包。pip install ogb")
完整的节点分类流水线
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import GCNConv, GATConv, SAGEConv
from torch_geometric.datasets import Planetoid
from torch_geometric.transforms import NormalizeFeatures
import matplotlib.pyplot as plt
# 加载数据
dataset = Planetoid(root='/tmp/Cora', name='Cora', transform=NormalizeFeatures())
data = dataset[0]
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
data = data.to(device)
class MultiLayerGNN(nn.Module):
"""组合多种 GNN 层的模型"""
def __init__(self, in_channels, hidden_channels, out_channels,
gnn_type="gcn", num_layers=3, dropout=0.5):
super().__init__()
self.dropout = dropout
self.gnn_type = gnn_type
self.convs = nn.ModuleList()
self.bns = nn.ModuleList()
# 输入层
self.convs.append(self._make_conv(in_channels, hidden_channels, gnn_type))
self.bns.append(nn.BatchNorm1d(hidden_channels))
# 中间层
for _ in range(num_layers - 2):
self.convs.append(self._make_conv(hidden_channels, hidden_channels, gnn_type))
self.bns.append(nn.BatchNorm1d(hidden_channels))
# 输出层
self.convs.append(self._make_conv(hidden_channels, out_channels, gnn_type))
def _make_conv(self, in_ch, out_ch, gnn_type):
if gnn_type == "gcn":
return GCNConv(in_ch, out_ch)
elif gnn_type == "sage":
return SAGEConv(in_ch, out_ch)
elif gnn_type == "gat":
return GATConv(in_ch, out_ch, heads=1)
else:
raise ValueError(f"Unknown GNN type: {gnn_type}")
def forward(self, x, edge_index):
for i, conv in enumerate(self.convs[:-1]):
x = conv(x, edge_index)
x = self.bns[i](x)
x = F.relu(x)
x = F.dropout(x, p=self.dropout, training=self.training)
return self.convs[-1](x, edge_index)
def run_experiment(gnn_type, epochs=200):
model = MultiLayerGNN(
in_channels=dataset.num_features,
hidden_channels=64,
out_channels=dataset.num_classes,
gnn_type=gnn_type,
num_layers=3
).to(device)
optimizer = torch.optim.Adam(
model.parameters(), lr=0.01, weight_decay=5e-4
)
train_losses = []
val_accs = []
for epoch in range(epochs):
# 训练
model.train()
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_losses.append(loss.item())
# 评估
model.eval()
with torch.no_grad():
out = model(data.x, data.edge_index)
pred = out.argmax(dim=1)
val_acc = pred[data.val_mask].eq(data.y[data.val_mask]).sum().item()
val_acc /= data.val_mask.sum().item()
val_accs.append(val_acc)
# 最终测试
model.eval()
with torch.no_grad():
out = model(data.x, data.edge_index)
pred = out.argmax(dim=1)
test_acc = pred[data.test_mask].eq(data.y[data.test_mask]).sum().item()
test_acc /= data.test_mask.sum().item()
return test_acc, train_losses, val_accs
# 比较不同的 GNN
results = {}
for gnn_type in ["gcn", "sage", "gat"]:
test_acc, losses, val_accs = run_experiment(gnn_type)
results[gnn_type] = test_acc
print(f"{gnn_type.upper():10s}: Test Accuracy = {test_acc:.4f}")
图分类完整示例
from torch_geometric.datasets import TUDataset
from torch_geometric.loader import DataLoader
from torch_geometric.nn import (
GINConv, global_mean_pool, global_add_pool
)
import torch
import torch.nn as nn
import torch.nn.functional as F
# 加载 MUTAG 数据集
dataset = TUDataset(root='/tmp/TUDataset', name='MUTAG')
dataset = dataset.shuffle()
# 划分训练/测试集
n = len(dataset)
train_dataset = dataset[:int(0.8 * n)]
test_dataset = dataset[int(0.8 * n):]
train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32)
class GIN(nn.Module):
"""
Graph Isomorphism Network (GIN) - 具有最大表达能力的 GNN
比 GCN 拥有更强的区分能力
"""
def __init__(self, in_channels, hidden_channels, out_channels,
num_layers=5, dropout=0.5):
super().__init__()
self.dropout = dropout
self.convs = nn.ModuleList()
self.bns = nn.ModuleList()
for i in range(num_layers):
in_ch = in_channels if i == 0 else hidden_channels
# GIN 的 MLP
mlp = nn.Sequential(
nn.Linear(in_ch, hidden_channels),
nn.BatchNorm1d(hidden_channels),
nn.ReLU(),
nn.Linear(hidden_channels, hidden_channels)
)
self.convs.append(GINConv(mlp, train_eps=True))
self.bns.append(nn.BatchNorm1d(hidden_channels))
# 图分类器
self.classifier = nn.Sequential(
nn.Linear(hidden_channels * num_layers, hidden_channels),
nn.ReLU(),
nn.Dropout(dropout),
nn.Linear(hidden_channels, out_channels)
)
def forward(self, x, edge_index, batch):
# 保存每一层的输出
xs = []
for conv, bn in zip(self.convs, self.bns):
x = conv(x, edge_index)
x = bn(x)
x = F.relu(x)
x = F.dropout(x, p=self.dropout, training=self.training)
xs.append(global_add_pool(x, batch)) # 图级别聚合
# 拼接所有层的图表示
out = torch.cat(xs, dim=1)
return self.classifier(out)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model_gin = GIN(
in_channels=dataset.num_features,
hidden_channels=64,
out_channels=dataset.num_classes
).to(device)
optimizer = torch.optim.Adam(model_gin.parameters(), lr=0.01)
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=50, gamma=0.5)
def train_gin():
model_gin.train()
total_loss = 0
for batch in train_loader:
batch = batch.to(device)
optimizer.zero_grad()
out = model_gin(batch.x, batch.edge_index, batch.batch)
loss = F.cross_entropy(out, batch.y)
loss.backward()
optimizer.step()
total_loss += loss.item()
return total_loss / len(train_loader)
def test_gin(loader):
model_gin.eval()
correct = 0
for batch in loader:
batch = batch.to(device)
with torch.no_grad():
pred = model_gin(batch.x, batch.edge_index, batch.batch).argmax(dim=1)
correct += pred.eq(batch.y).sum().item()
return correct / len(loader.dataset)
for epoch in range(1, 201):
loss = train_gin()
train_acc = test_gin(train_loader)
test_acc = test_gin(test_loader)
scheduler.step()
if epoch % 20 == 0:
print(f"Epoch {epoch:03d} | Loss: {loss:.4f} | "
f"Train: {train_acc:.4f} | Test: {test_acc:.4f}")
8. DGL (Deep Graph Library) 对比
# DGL 示例 - 与 PyG 对比
# pip install dgl
try:
import dgl
import dgl.nn as dglnn
import torch
import torch.nn as nn
import torch.nn.functional as F
class DGLGCN(nn.Module):
"""使用 DGL 实现的 GCN"""
def __init__(self, in_feats, hidden_size, num_classes):
super().__init__()
self.conv1 = dglnn.GraphConv(in_feats, hidden_size)
self.conv2 = dglnn.GraphConv(hidden_size, num_classes)
def forward(self, g, features):
x = F.relu(self.conv1(g, features))
x = F.dropout(x, training=self.training)
return self.conv2(g, x)
# 创建 DGL 图
src = torch.tensor([0, 1, 2, 3, 4])
dst = torch.tensor([1, 2, 3, 4, 0])
g = dgl.graph((src, dst))
g.ndata['feat'] = torch.randn(5, 16)
model_dgl = DGLGCN(16, 32, 4)
out = model_dgl(g, g.ndata['feat'])
print(f"DGL GCN output: {out.shape}")
except ImportError:
print("未安装 DGL。pip install dgl")
PyG 与 DGL 对比：
| 特性 | PyTorch Geometric (PyG) | Deep Graph Library (DGL) |
|---|---|---|
| API 风格 | PyTorch 原生 | 框架无关 |
| 数据表示 | edge_index (COO) | DGLGraph 对象 |
| 速度 | 非常快 | 快 |
| 社区 | 大规模 | 大规模 |
| 模型数量 | 非常多 | 多 |
| 学习曲线 | 低 | 中等 |
9. 实战应用
分子性质预测 (OGB)
try:
from ogb.graphproppred import PygGraphPropPredDataset, Evaluator
from torch_geometric.loader import DataLoader
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import GINEConv, global_mean_pool
# 加载 HIV 分子数据集
dataset_mol = PygGraphPropPredDataset(name='ogbg-molhiv')
split_idx = dataset_mol.get_idx_split()
train_loader_mol = DataLoader(
dataset_mol[split_idx["train"]],
batch_size=32,
shuffle=True
)
class MoleculeGNN(nn.Module):
"""分子性质预测模型"""
def __init__(self, hidden_channels=300, num_layers=5):
super().__init__()
self.atom_encoder = nn.Embedding(100, hidden_channels)
self.bond_encoder = nn.Embedding(10, hidden_channels)
self.convs = nn.ModuleList()
for _ in range(num_layers):
mlp = nn.Sequential(
nn.Linear(hidden_channels, hidden_channels * 2),
nn.BatchNorm1d(hidden_channels * 2),
nn.ReLU(),
nn.Linear(hidden_channels * 2, hidden_channels)
)
self.convs.append(GINEConv(mlp))
self.pool = global_mean_pool
self.predictor = nn.Linear(hidden_channels, 1)
def forward(self, x, edge_index, edge_attr, batch):
x = self.atom_encoder(x.squeeze())
edge_attr = self.bond_encoder(edge_attr.squeeze())
for conv in self.convs:
x = conv(x, edge_index, edge_attr)
x = F.relu(x)
graph_embed = self.pool(x, batch)
return self.predictor(graph_embed)
print("OGB 分子数据集加载成功")
except ImportError:
print("未安装 ogb 包。pip install ogb")
推荐系统
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import LightGCN
class RecommendationSystem(nn.Module):
"""
基于 LightGCN 的协同过滤
在用户-物品二分图上学习嵌入
"""
def __init__(self, num_users, num_items, embedding_dim=64, num_layers=3):
super().__init__()
self.num_users = num_users
self.num_items = num_items
self.embedding_dim = embedding_dim
self.num_layers = num_layers
# 用户/物品嵌入
self.user_emb = nn.Embedding(num_users, embedding_dim)
self.item_emb = nn.Embedding(num_items, embedding_dim)
# LightGCN: 不做非线性变换, 只做简单聚合
self.lightgcn = LightGCN(
num_nodes=num_users + num_items,
embedding_dim=embedding_dim,
num_layers=num_layers
)
self._init_weights()
def _init_weights(self):
nn.init.normal_(self.user_emb.weight, std=0.01)
nn.init.normal_(self.item_emb.weight, std=0.01)
def forward(self, edge_index):
# 全部节点嵌入
x = torch.cat([self.user_emb.weight, self.item_emb.weight], dim=0)
# LightGCN 传播
embeddings = self.lightgcn(x, edge_index)
return embeddings[:self.num_users], embeddings[self.num_users:]
def predict(self, user_ids, item_ids, edge_index):
user_embs, item_embs = self(edge_index)
u = user_embs[user_ids]
i = item_embs[item_ids]
return (u * i).sum(dim=1)
# BPR 损失函数
def bpr_loss(pos_scores, neg_scores):
"""Bayesian Personalized Ranking Loss"""
return -F.logsigmoid(pos_scores - neg_scores).mean()
10. 图生成模型
GraphVAE
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch_geometric.nn import GCNConv
class GraphVAE(nn.Module):
"""图变分自编码器"""
def __init__(self, in_channels, hidden_channels, latent_dim):
super().__init__()
# 编码器 (GNN)
self.conv1 = GCNConv(in_channels, hidden_channels)
self.conv_mu = GCNConv(hidden_channels, latent_dim)
self.conv_logvar = GCNConv(hidden_channels, latent_dim)
def encode(self, x, edge_index):
h = F.relu(self.conv1(x, edge_index))
mu = self.conv_mu(h, edge_index)
logvar = self.conv_logvar(h, edge_index)
return mu, logvar
def reparameterize(self, mu, logvar):
if self.training:
std = torch.exp(0.5 * logvar)
eps = torch.randn_like(std)
return mu + eps * std
return mu
def decode(self, z):
# 通过内积计算边的概率
adj_pred = torch.sigmoid(z @ z.t())
return adj_pred
def forward(self, x, edge_index):
mu, logvar = self.encode(x, edge_index)
z = self.reparameterize(mu, logvar)
adj_pred = self.decode(z)
return adj_pred, mu, logvar
def loss(self, adj_pred, adj_target, mu, logvar):
# 重构损失
recon_loss = F.binary_cross_entropy(adj_pred, adj_target)
# KL 散度
kl_loss = -0.5 * torch.mean(
1 + logvar - mu.pow(2) - logvar.exp()
)
return recon_loss + kl_loss
测验
Q1. GCN 与 GAT 最大的区别是什么？
答案：GCN 用固定权重（基于度数的归一化）聚合所有邻居节点，而 GAT 通过注意力机制为每个邻居动态学习不同的权重。
解析：GCN 的聚合权重由节点的度数（degree）固定决定。而 GAT 的注意力系数基于相连两个节点的特征向量动态计算，因此能让模型对更重要的邻居给予更多关注。用多头注意力提升稳定性，也是 GAT 的一个优点。
Q2. GraphSAGE 相比 GCN 在归纳式学习上更有优势的原因是什么？
答案：因为 GraphSAGE 学习的是聚合函数（aggregator），可以用它为新节点的邻居生成嵌入。
解析：GCN 在训练时需要用到整张图的邻接矩阵，一旦有新节点加入就得重新训练，属于直推式（transductive）方法。GraphSAGE 学习的是采样并聚合邻居的函数，因此即便是训练时从未见过的新节点，也可以套用这个函数来生成嵌入。这一特性在 Pinterest、LinkedIn 等动态变化的大规模图中被实际应用。
Q3. 消息传递（MPNN）框架的三个阶段是什么？
答案：Message（消息计算）、Aggregate（聚合）、Update（更新）三个阶段。
解析：Message 阶段针对每条边计算要从邻居节点传递过来的消息。Aggregate 阶段将节点收到的所有邻居消息通过求和、平均、取最大值等方式聚合起来。Update 阶段将聚合后的消息与当前节点的嵌入结合，生成新的节点嵌入。GCN、GAT、GraphSAGE、GIN 等大多数 GNN 都可以纳入这一框架统一描述。
Q4. 过平滑（Over-smoothing）问题是什么？如何解决？
答案：随着层数加深，所有节点的嵌入趋于相似的现象。可以用残差连接、JK-Net、DropEdge 等方法缓解。
解析：K 层 GNN 会聚合 K 跳邻居的信息。层数越多，覆盖的邻居范围就越广，最终所有节点都会收敛到同一个全局平均值。残差连接（Residual connections）通过直接传递前一层的信息来保留节点的独有信息。JK-Net（Jumping Knowledge Networks）在最终表示中利用所有层的嵌入。DropEdge 在训练时随机移除部分边。
Q5. GNN 的表达能力与 WL 测试等价，这意味着什么？
答案：标准 GNN 会将 Weisfeiler-Leman（WL）图同构测试也无法区分的两张图，嵌入为相同的表示。
解析：WL 测试是一种通过反复聚合并哈希邻居标签来判断两张图是否同构（isomorphic）的算法。Xu 等人（2019）通过 GIN（Graph Isomorphism Network）证明了标准 GNN 的表达能力至多与 1-WL 测试相当。在 WL 测试无法区分的图对上，GNN 同样无法区分。为了突破这一局限，目前正在研究对应更强的 k 阶 WL 测试的高阶 GNN。
结语
本指南梳理了图神经网络的完整生态：
- 图论基础：节点、边、邻接矩阵、图的特性
- 消息传递范式：GNN 的核心原理
- 主要架构：GCN、GraphSAGE、GAT、Graph Transformer、GIN
- 图级别预测：Global Pooling、DiffPool
- 链接预测：知识图谱、推荐系统
- PyTorch Geometric：节点分类、图分类完整示例
- 实战应用：分子设计、推荐系统、欺诈检测
- 图生成模型：GraphVAE
GNN 正在分子设计、药物发现、社交网络分析、交通预测、推荐系统等众多领域取得创新性的成果。得益于 PyTorch Geometric、DGL 等库，实现变得越来越简单，OGB 等基准测试也让公平比较成为可能。