别再死记公式了！用Python代码手搓一个Graph Transformer，直观理解它与GNN/Transformer的异同

用Python手搓Graph Transformer：从代码透视GNN与Transformer的融合奥秘

在深度学习领域，图神经网络(GNN)和Transformer架构如同两颗璀璨的明珠，分别照亮了非欧式空间数据与序列建模的疆域。而当这两大范式相遇，便催生出了Graph Transformer这一充满潜力的混合体。本文将带您从零开始，用PyTorch实现一个精简版的Graph Transformer，通过可运行的代码示例，直观展示它与传统GNN、标准Transformer的核心差异。

1. 环境准备与数据加载

首先确保您的Python环境已安装以下关键库：

pip install torch torch-geometric numpy matplotlib

我们将使用Cora引文网络数据集作为演示对象，这是一个经典的图机器学习基准数据集：

import torch from torch_geometric.datasets import Planetoid from torch_geometric.utils import to_dense_adj dataset = Planetoid(root='/tmp/Cora', name='Cora') data = dataset[0] # 将稀疏邻接矩阵转换为稠密形式 adj = to_dense_adj(data.edge_index)[0] num_nodes = data.num_nodes feat_dim = dataset.num_features

关键参数说明：

• data.x: 节点特征矩阵 (2708×1433)
• adj: 邻接矩阵 (2708×2708)
• data.edge_index: 边信息的稀疏表示

2. 三大模型架构对比实现

2.1 传统GNN层实现

典型的GNN通过聚合邻居信息来更新节点表示：

import torch.nn as nn import torch.nn.functional as F class GNNLayer(nn.Module): def __init__(self, in_dim, out_dim): super().__init__() self.linear = nn.Linear(in_dim, out_dim) def forward(self, x, adj): # 邻居信息聚合 neighbor_agg = torch.matmul(adj, x) # 结合自身信息 h = self.linear(neighbor_agg) return F.relu(h)

核心特点：

• 仅考虑直接邻居节点（1-hop）
• 计算复杂度与边数量线性相关
• 天然保留图拓扑结构

2.2 标准Transformer层实现

原始Transformer的注意力机制不考虑图结构：

class TransformerLayer(nn.Module): def __init__(self, dim, heads=4): super().__init__() self.attention = nn.MultiheadAttention(dim, heads) self.norm = nn.LayerNorm(dim) def forward(self, x): # 全局注意力计算 attn_out, _ = self.attention(x, x, x) return self.norm(attn_out + x)

关键差异：

• 计算所有节点间的注意力权重
• 完全忽略原始图结构
• 计算复杂度与节点数平方成正比

2.3 Graph Transformer层实现

融合二者优势的Graph Transformer实现：

class GraphTransformerLayer(nn.Module): def __init__(self, dim, heads=4): super().__init__() self.attention = nn.MultiheadAttention(dim, heads) self.norm = nn.LayerNorm(dim) self.edge_encoder = nn.Linear(1, heads) # 边特征编码 def forward(self, x, adj): # 生成注意力偏置 (考虑图结构) edge_embed = self.edge_encoder(adj.unsqueeze(-1)) attn_bias = edge_embed.permute(2,0,1) # [heads, N, N] # 带结构偏置的注意力计算 attn_out, _ = self.attention(x, x, x, attn_mask=attn_bias) return self.norm(attn_out + x)

创新点对比：

3. 关键组件深度解析

3.1 注意力机制改造

Graph Transformer的核心创新在于对注意力矩阵的结构化约束：

def scaled_dot_product_attention(Q, K, V, adj_mask=None): scores = torch.matmul(Q, K.transpose(-2,-1)) / torch.sqrt(dim) if adj_mask is not None: scores = scores + adj_mask # 结构偏置注入 weights = F.softmax(scores, dim=-1) return torch.matmul(weights, V)

结构偏置的常见实现方式：

1. 二进制邻接矩阵掩码
2. 基于节点距离的衰减系数
3. 可学习的边编码器

3.2 位置编码的图适配

传统Transformer的位置编码(PE)在图数据中的改进：

class GraphPositionalEncoding(nn.Module): def __init__(self, dim, max_len=100): super().__init__() # 基于节点中心度的编码 self.centrality_encoder = nn.Linear(1, dim//2) # 基于随机游走的编码 self.rw_encoder = nn.Linear(max_len, dim//2) def forward(self, x, degree, rw_pos): cent_enc = self.centrality_encoder(degree.unsqueeze(-1)) rw_enc = self.rw_encoder(rw_pos) return x + torch.cat([cent_enc, rw_enc], dim=-1)

4. 完整模型训练实战

构建一个端到端的Graph Transformer分类模型：

class GraphTransformer(nn.Module): def __init__(self, in_dim, hidden_dim, out_dim, heads=4): super().__init__() self.embed = nn.Linear(in_dim, hidden_dim) self.layers = nn.ModuleList([ GraphTransformerLayer(hidden_dim, heads) for _ in range(3) ]) self.classifier = nn.Linear(hidden_dim, out_dim) def forward(self, x, adj): h = self.embed(x) for layer in self.layers: h = layer(h, adj) return F.log_softmax(self.classifier(h), dim=-1) # 训练循环示例 model = GraphTransformer(feat_dim, 64, dataset.num_classes) optimizer = torch.optim.Adam(model.parameters(), lr=0.01) for epoch in range(200): model.train() optimizer.zero_grad() out = model(data.x, adj) loss = F.nll_loss(out[data.train_mask], data.y[data.train_mask]) loss.backward() optimizer.step()

训练技巧：

• 使用梯度裁剪防止爆炸
• 添加层归一化稳定训练
• 采用标签平滑提升泛化

5. 效果对比与可视化分析

在Cora数据集上的性能对比：

注意力模式可视化展示：

import matplotlib.pyplot as plt def plot_attention(head_weights, node_idx): plt.figure(figsize=(10,5)) for i, weights in enumerate(head_weights): plt.subplot(1, len(head_weights), i+1) plt.imshow(weights[node_idx].detach().numpy()) plt.title(f'Head {i+1}') plt.show()

通过可视化可见：

• 某些注意力头自动聚焦局部邻居
• 部分头捕获长距离依赖关系
• 边编码有效保留了结构信息