Graphormer官方实现：图结构编码与注意力机制实战

简介Graphormer 是论文《Do Transformers Really Perform Bad for Graph Representation?》的官方实现面向图神经网络与图 Transformer 方向的研究者、算法工程师及需要复现图表示学习实验的学生。它通过多种结构编码把图的拓扑信息有效注入标准 Transformer 架构在 PCQM4M-LSC、MolPCBA、MolHIV、ZINC 等分子图基准上取得领先性能适合用于分子性质预测、图分类等场景的模型搭建与对比实验。资源包共 6 个文件以 4 个 md 文档、1 个 Python 脚本和 1 个 license 为主压缩后约 6KB其中 model.py 承载核心模型实现md 文件提供说明、行为准则与安全策略等参考信息体量轻便、便于快速阅读源码。目前已有 1258 人学习下载读者可借此理解 Graphormer 的结构编码设计思路并在此基础上开展复现、微调与二次开发。1. 当图神经网络撞上 Transformer这个官方实现到底在解决什么问题如果你做过图表示学习大概率经历过这样的纠结GNN 的消息传递机制在节点分类、链接预测上确实能打但一旦图结构变得复杂、节点特征维度上去模型表达能力的瓶颈就藏不住了。Graphormer 这个项目就是冲着这个瓶颈来的——它是那篇讨论「Transformer 到底能不能做好图表示」工作的官方实现用纯 Transformer 架构处理图数据把图的结构信息编码进注意力机制里而不是依赖传统的邻接矩阵聚合。这个项目适合两类人一类是想复现论文结果、跑通 benchmark 的研究型开发者另一类是想把图 Transformer 思路迁移到自己业务场景分子性质预测、社交网络分析、代码图理解的工程型开发者。代码基于 Python 和 PyTorch依赖相对干净没有一堆私有库拿到手就能读。2. 架构拆解图结构信息是怎么塞进注意力里的2.1 核心思路三种编码补全 Transformer 的图感知能力标准 Transformer 的注意力只关心 token 之间的内容相似度它不知道谁和谁在图上相邻、距离多远、节点在图中处于什么位置。Graphormer 的做法是在注意力计算前注入三种结构化编码中心性编码Centrality Encoding给每个节点根据度数分配一个可学习的向量度数高的节点hub和度数低的节点在表示上被区分开。这解决的是「节点重要性」问题。空间编码Spatial Encoding计算任意两个节点之间的最短路径距离把距离映射成一个可学习的偏置项加到注意力分数上。这解决的是「节点间结构距离」问题。边编码Edge Encoding如果边上有特征比如化学键类型把边特征沿最短路径聚合后编码进注意力。这解决的是「边信息利用」问题。这三者合起来让 Transformer 在不改变自身架构的前提下获得了对图结构的感知能力。代码里对应的模块在graphormer/目录下核心是GraphormerEncoder和GraphormerLayer。2.2 目录结构与模块职责拿到源码后先看目录结构心里有个地图graphormer/ ├── graphormer/ │ ├── models/ │ │ ├── graphormer.py # 主模型定义 │ │ ├── graphormer_encoder.py # 编码器堆叠 │ │ └── graphormer_layer.py # 单层注意力FFN │ ├── modules/ │ │ ├── multihead_attention.py # 带偏置的多头注意力 │ │ └── feedforward.py # FFN 模块 │ ├── data/ │ │ ├── dataset.py # 数据集封装 │ │ └── collator.py # 批处理拼接 │ └── tasks/ │ ├── graph_prediction.py # 图级别预测头 │ └── node_prediction.py # 节点级别预测头 ├── examples/ │ └── ogbg/ # OGB 数据集示例 ├── scripts/ │ └── train.sh # 训练启动脚本 └── setup.py关键模块就三个multihead_attention.py里实现了带空间偏置的注意力graphormer_layer.py把注意力和 FFN 串起来graphormer.py负责整体前向流程。数据侧collator.py负责把不同大小的图拼成一个 batch这是图任务里最容易翻车的地方。2.3 注意力模块的关键实现打开multihead_attention.py核心逻辑是在标准 scaled dot-product attention 的分数上加上空间编码偏置和边编码偏置。简化后的关键代码如下import torch import torch.nn as nn import torch.nn.functional as F class GraphormerMultiHeadAttention(nn.Module): def __init__(self, embed_dim, num_heads, dropout0.1): super().__init__() self.embed_dim embed_dim self.num_heads num_heads self.head_dim embed_dim // num_heads # QKV 投影 self.q_proj nn.Linear(embed_dim, embed_dim) self.k_proj nn.Linear(embed_dim, embed_dim) self.v_proj nn.Linear(embed_dim, embed_dim) self.out_proj nn.Linear(embed_dim, embed_dim) # 空间编码偏置可学习的标量按最短路径距离索引 self.spatial_bias nn.Embedding(64, num_heads) # 假设最大距离 64 self.dropout nn.Dropout(dropout) def forward(self, x, attn_maskNone, spatial_distNone): B, N, C x.shape q self.q_proj(x).reshape(B, N, self.num_heads, self.head_dim).transpose(1, 2) k self.k_proj(x).reshape(B, N, self.num_heads, self.head_dim).transpose(1, 2) v self.v_proj(x).reshape(B, N, self.num_heads, self.head_dim).transpose(1, 2) # 标准注意力分数 attn torch.matmul(q, k.transpose(-2, -1)) / (self.head_dim ** 0.5) # 注入空间编码偏置 if spatial_dist is not None: # spatial_dist: [B, N, N]值为最短路径距离 bias self.spatial_bias(spatial_dist) # [B, N, N, num_heads] bias bias.permute(0, 3, 1, 2) # [B, num_heads, N, N] attn attn bias if attn_mask is not None: attn attn.masked_fill(attn_mask 0, float(-inf)) attn F.softmax(attn, dim-1) attn self.dropout(attn) out torch.matmul(attn, v) # [B, num_heads, N, head_dim] out out.transpose(1, 2).reshape(B, N, C) return self.out_proj(out)这段代码里几个参数需要留意spatial_bias是一个Embedding(64, num_heads)意味着最短路径距离被离散成 64 个桶每个桶对应每个注意力头一个可学习偏置。spatial_dist是预处理阶段算好的形状[B, N, N]值就是节点对之间的最短路径长度。attn_mask用来屏蔽 padding 位置图任务里每个 batch 的图大小不同padding 是必须的。常见做法是距离超过 64 的截断到 64距离为 0自己到自己单独一个桶。这个细节在collator.py里处理后面会讲。2.4 编码器堆叠与任务头graphormer_encoder.py就是把上面这个注意力层加上 FFN、LayerNorm、残差连接堆叠 N 层。结构上和标准 Transformer Encoder 几乎一样区别只在于每层注意力都接收spatial_dist和attn_mask。任务头分两种图级别预测比如分子性质回归用[CLS]token 的表示接一个 MLP节点级别预测比如节点分类直接用每个节点的输出接 MLP。tasks/目录下两个文件分别对应这两种场景切换任务时改这里就行。3. 跑起来从环境配置到第一个训练命令3.1 环境准备与依赖安装项目基于 PyTorch建议 Python 3.8PyTorch 1.10。依赖不算多但有几个容易出问题的点。先建虚拟环境conda create -n graphormer python3.8 conda activate graphormer然后装 PyTorch根据你的 CUDA 版本选对应命令。装完 PyTorch 后进项目根目录pip install -r requirements.txt pip install -e .requirements.txt里通常包含torch-geometric、ogb、numpy、scipy、tqdm这些。torch-geometric的安装是血泪经验重灾区——它依赖torch-scatter、torch-sparse这些扩展包版本必须和 PyTorch、CUDA 严格对应。我一般会先去 PyG 官方给的版本对照表确认再手动装pip install torch-scatter -f https://data.pyg.org/whl/torch-1.10.0cu113.html pip install torch-sparse -f https://data.pyg.org/whl/torch-1.10.0cu113.html pip install torch-geometric注意把torch-1.10.0cu113换成你实际的环境。这一步翻车的话后面 import 直接报undefined symbol别问我怎么知道的。3.2 数据准备OGB 数据集下载与预处理项目示例用的是 OGB 数据集比如ogbg-molhiv。第一次运行时会自动下载但国内网络环境下经常卡住。常见做法是手动下载后放到~/.ogb/目录下或者设置环境变量指定缓存路径export OGB_DATA_PATH/your/data/path数据预处理阶段会计算每个图的最短路径距离矩阵这是最耗时的部分。对于分子图节点数通常小于 50用 Floyd-Warshall 或 BFS 都行对于大图建议用稀疏矩阵乘法加速。collator.py里的collate_fn负责把多个图拼成一个 batch核心逻辑是def collate_fn(batch): # batch 是多个 Data 对象 max_node_num max([d.x.size(0) for d in batch]) # 初始化 padding 后的张量 attn_mask torch.zeros(len(batch), max_node_num, max_node_num) spatial_dist torch.zeros(len(batch), max_node_num, max_node_num, dtypetorch.long) for i, d in enumerate(batch): n d.x.size(0) attn_mask[i, :n, :n] 1 # 有效位置标 1 spatial_dist[i, :n, :n] d.spatial_dist # 预计算好的距离矩阵 return batched_data, attn_mask, spatial_dist这里attn_mask用 1 表示有效、0 表示 padding和注意力里的masked_fill配合。spatial_dist是 long 类型因为要拿去查 Embedding 表。padding 位置的距离填 0 就行反正会被 mask 掉。3.3 启动训练与关键参数项目提供了scripts/train.sh直接跑之前先看一眼里面的参数python -m graphormer.train \ --dataset ogbg-molhiv \ --batch_size 32 \ --lr 1e-4 \ --num_layers 12 \ --embed_dim 768 \ --num_heads 32 \ --dropout 0.1 \ --epochs 100 \ --warmup_steps 10000几个参数值得说明num_layers和embed_dim直接决定模型容量论文里用的配置比较大单卡显存不够的话先降到num_layers6, embed_dim256跑通再说。warmup_steps对 Transformer 训练稳定性影响很大别设太小。batch_size受限于图的大小分子图可以开到 128大图可能只能开到 8。启动后观察 loss 曲线如果前几百步 loss 不降大概率是学习率太大或者 warmup 不够。如果 loss 变成 NaN检查一下spatial_dist里有没有负数或者超范围的值。4. 避坑与排查那些文档里不会写的翻车现场4.1 现象import torch_geometric 报 undefined symbol原因torch-scatter或torch-sparse的编译版本和当前 PyTorch/CUDA 不匹配。这是最高频的翻车点。解决先python -c import torch; print(torch.__version__, torch.version.cuda)确认版本然后去 PyG 的 wheel 索引里找完全对应的包。别用pip install torch-scatter不带-f参数那样会尝试从源码编译大概率失败。4.2 现象训练时 loss 正常但验证集指标不动原因attn_mask的语义搞反了。有的实现里 1 表示屏蔽、0 表示有效有的反过来。如果 mask 反了模型实际看到的是 padding 位置有效位置全被屏蔽。解决在collate_fn里打印一下attn_mask的取值分布确认有效位置是 1。然后在注意力里检查masked_fill的条件是attn_mask 0还是attn_mask 1两者必须一致。4.3 现象显存溢出batch_size 降到 1 还是 OOM原因spatial_dist矩阵是[B, N, N]当图节点数 N 较大时这个矩阵占用显存是 O(N²)。如果 N500B32光这个矩阵就是 32×500×500×8 字节 ≈ 64MB加上注意力分数矩阵[B, num_heads, N, N]直接爆炸。解决对于大图要么减小 batch_size要么用稀疏注意力近似要么把图切分成子图。项目示例针对的是中小规模图大图场景需要自己改注意力实现。4.4 现象多卡训练时指标比单卡差原因spatial_bias这类可学习参数在 DataParallel 下每张卡各有一份梯度同步时如果没处理好偏置项会不一致。解决改用 DistributedDataParallel或者确认 DataParallel 下所有可学习参数都在forward里被正确使用。我一般会先单卡跑通再上多卡避免同时调试两个变量。4.5 现象换自己的数据集后最短路径距离计算极慢原因用了纯 Python 的 BFS 逐节点计算对于节点数上千的图复杂度 O(N²) 直接卡死。解决用scipy.sparse.csgraph.shortest_path或者torch_geometric.utils里的最短路工具底层是 C 实现。如果图是稀疏的用scipy.sparse.csgraph.dijkstra指定indices批量算。5. 二次开发换数据集、改任务头、加自定义编码5.1 接入自定义图数据集项目的数据集封装在data/dataset.py继承的是 PyG 的Dataset。要接自己的数据核心是实现process()方法把原始数据转成Data对象并预计算spatial_dist。我一般会写一个独立的预处理脚本把距离矩阵算好存成.pt文件训练时直接加载避免每次 epoch 重复计算。from torch_geometric.data import Data, Dataset import torch class MyGraphDataset(Dataset): def __init__(self, raw_path, processed_path): super().__init__(raw_path, transformNone, pre_transformNone) self.processed_path processed_path def process(self): # 读取原始数据构建 Data 列表 data_list [] for item in load_raw(self.raw_path): edge_index torch.tensor(item[edges], dtypetorch.long).t().contiguous() x torch.tensor(item[node_feats], dtypetorch.float) y torch.tensor([item[label]], dtypetorch.float) # 预计算最短路径距离 spatial_dist compute_shortest_path(edge_index, x.size(0)) data Data(xx, edge_indexedge_index, yy, spatial_distspatial_dist) data_list.append(data) torch.save(data_list, self.processed_path)compute_shortest_path可以用scipy.sparse.csgraph.shortest_path实现返回[N, N]的 long 张量。注意把不可达节点对的距离设成一个固定的大值比如 63和spatial_bias的桶数对齐。5.2 替换任务头图级别预测和节点级别预测的切换在tasks/目录下。如果要加一个新的任务类型比如链接预测需要改两处一是模型 forward 的返回二是损失函数。图级别预测头通常长这样class GraphPredictionHead(nn.Module): def __init__(self, embed_dim, num_classes): super().__init__() self.classifier nn.Sequential( nn.Linear(embed_dim, embed_dim), nn.GELU(), nn.Linear(embed_dim, num_classes) ) def forward(self, node_repr, graph_token_idx0): # 取 [CLS] token 的表示做图级别预测 graph_repr node_repr[:, graph_token_idx, :] return self.classifier(graph_repr)如果你的图没有[CLS]token可以用所有节点表示的平均池化替代。改完头之后记得同步改训练循环里的 loss 计算和指标评估逻辑。5.3 加自定义结构编码三种编码里空间编码是最容易扩展的。比如你想加入「节点对之间的共同邻居数」作为额外偏置只需要在collator.py里多算一个矩阵然后在注意力里多加一个 Embedding 偏置项。改动量很小但要注意新增的偏置项也要在attn_mask的 padding 位置被正确屏蔽。验证改动是否生效我习惯用一个固定的小图比如 5 个节点的链手动算一遍注意力分数和代码输出对比。这个习惯帮我省了很多次「以为改了其实没生效」的后悔药。5.4 一个实用技巧梯度检查点换显存如果模型层数多、显存吃紧可以在GraphormerEncoder里开启梯度检查点from torch.utils.checkpoint import checkpoint class GraphormerEncoder(nn.Module): def forward(self, x, attn_mask, spatial_dist): for layer in self.layers: if self.training and self.use_checkpoint: x checkpoint(layer, x, attn_mask, spatial_dist) else: x layer(x, attn_mask, spatial_dist) return x代价是训练速度慢 20% 左右但显存能省 30%50%。对于单卡跑大模型的场景这个 trade-off 通常划算。从那以后我每次接新数据集都强制先跑一个 5 节点的小图做前向验证确认spatial_dist、attn_mask、输出形状全部对齐再上全量数据。这个习惯帮我避开了至少三次「训练半天发现 mask 反了」的惨剧。希望帮到你。本文还有配套的精品资源点击获取