简介Graphormer 官方实现代码包面向图神经网络与 Transformer 方向的研究者、Python 开发者及需要复现图表示学习实验的学生。它围绕“Transformer 能否有效处理图结构”这一问题提供将图结构信息编码进标准 Transformer 的核心实现可用于论文复现、模型对照和二次开发。压缩包共 6 个文件由 4 个 Markdown 文档、1 个 Python 脚本和 1 份许可证组成整体约 6KB体量轻便便于快速浏览代码与说明。已有 1258 人学习下载具备一定参考热度。读者可从中获取模型主体代码、项目说明与授权信息理解结构编码的设计思路及实验配置为复现 PCQM4M-LSC、MolPCBA 等基准结果或迁移到自己的图数据任务提供起点。1. 图 Transformer 不是玄学Graphormer 把结构偏置塞进注意力哪一层一个做分子性质预测的团队GNN 调了两个月验证集指标卡住不动换 Graphormer 后第一版配置就抬了一截。反直觉的地方在于普通 Transformer 直接套图数据往往翻车不是注意力机制不行而是它看不见节点度数、最短路径和边类型。Graphormer 的官方 Python 实现把这三类结构信号编码成注意力偏置让图 Transformer 真正吃到拓扑信息。这份资源适合已有 PyTorch 基础、想在图级任务上试 Transformer 的从业者不适合只想调包、不愿做图预处理的人。它解决的是从论文公式到可跑代码之间的断层但代价是显存、数据格式和参数耦合都比普通 GNN 更敏感。下面按“编码原理、环境跑通、参数设置、踩坑排查、迁移微调”的顺序拆。2. 拆开 Graphormer 的三类编码普通 Transformer 在图上翻车的三个位置普通 Transformer 的注意力是置换不变的把节点打乱顺序结果不变。图数据恰恰需要顺序之外的结构信息哪个节点是枢纽、两个节点隔了几条边、边是什么类型。Graphormer 没有改注意力公式而是在注意力 logits 上加偏置并把结构信号塞进节点表示。理解这三类编码后面调参和排错才有依据。2.1 节点中心性编码把“谁更重要”塞进注意力图里节点的度分布通常长尾少数节点连接很多邻居多数节点只有两三个邻居。普通 Transformer 的节点嵌入只来自特征不知道谁在结构上更关键。Graphormer 用入度和出度的可学习嵌入来补这个信号常见写法是把中心性嵌入加到节点特征上或者作为注意力偏置的一部分。对无向图入度和出度相同但实现里仍常保留两套嵌入方便有向图复用。中心性编码的关键参数是嵌入表大小和截断阈值。度数是整数理论上可以到几千但绝大多数节点度数很小。如果直接按最大度数建嵌入表表会很大训练时大量 id 很少被激活。常见做法是设一个上限超过上限的度数共享同一个 id相当于把长尾截断。这个上限不是真实度数上限而是嵌入表容量设小了会损失高 degree 节点的区分度设大了浪费显存。import torch def degree_ids(edge_index, num_nodes, max_degree512): # edge_index: [2, E]有向边无向图正反各存一次 src, dst edge_index[0], edge_index[1] in_deg torch.bincount(dst, minlengthnum_nodes) out_deg torch.bincount(src, minlengthnum_nodes) # 长尾截断超过 max_degree 的节点共享最大 id in_deg torch.clamp(in_deg, maxmax_degree) out_deg torch.clamp(out_deg, maxmax_degree) return in_deg, out_deg这段代码只做一件事把边索引转成每个节点的入度、出度 id。max_degree直接决定中心性嵌入表的行数我一般先从 512 试起图特别大再降到 256。num_nodes要和节点特征行数一致否则后面索引会越界。无向图必须把正反边都写进edge_index只存单向边会让入度出度算错模型仍能跑但中心性信号是歪的。还有一个容易忽略的点中心性编码不是万能的。如果任务本身跟度数关系弱比如某些节点分类任务中心性嵌入可能带来过拟合。官方实现把它做成可配置项调参时可以先关掉对比一下验证集。我的习惯是先用默认中心性跑一版再关掉跑一版差值超过两个点才保留。2.2 空间编码与边编码让注意力看见路径和边特征普通注意力只算 Query 和 Key 的点积两个节点特征相似注意力就高但它们在图里可能隔了很远。Graphormer 在注意力 logits 上加一个偏置项这个偏置来自两节点之间的最短路径距离和路径上的边类型。空间编码负责“隔多远”边编码负责“路上经过什么类型的边”。空间编码通常是一张距离查找表先算所有节点对的最短路径距离超过阈值就截断到最大值。阈值设多少等于空间嵌入表有多大。分子图节点数不大阈值 20 左右常见通用图节点多阈值要小否则表膨胀。边编码更细沿着最短路径取边特征序列用一维卷积或嵌入把序列压成一个向量再映射成节点对的偏置。这样两个节点即使特征相似只要路径上的边类型不同注意力也会被拉开。from collections import deque import torch def bfs_spatial_pos(adj, num_nodes, max_dist20): # adj: list[list[int]]无权无向邻接表 spatial torch.full((num_nodes, num_nodes), max_dist, dtypetorch.long) for s in range(num_nodes): dist [-1] * num_nodes dist[s] 0 q deque([s]) while q: u q.popleft() if dist[u] max_dist: continue for v in adj[u]: if dist[v] -1: dist[v] dist[u] 1 q.append(v) for t, d in enumerate(dist): if d ! -1: spatial[s, t] min(d, max_dist) return spatialbfs_spatial_pos输出的是节点对最短路径距离矩阵max_dist既是截断阈值也是空间嵌入表的行数。超过阈值的节点对统一落到最大 id不再区分 21 跳还是 30 跳。这个函数复杂度是 O(N*(NE))节点数上千时就会明显拖慢预处理。我的经验是节点数超过 500 的图不要在全图上直接算先做子图采样或者把任务改成节点级局部窗口。边编码的输入是路径上的边类型序列。分子图里边类型少键型、芳香性、是否成环几类就够通用图可能要按边特征分桶。桶数就是边编码嵌入表的行数设少了不同边类型挤在一起设多了稀疏。路径长度也要截断常见上限 5 到 10超过上限的路径只保留前几跳。这里没有后悔药路径长度和边类型桶数一旦定错后面调模型很难救回来只能重新预处理。2.3 从公式到张量编码如何落到 PyTorch 模块里把三类编码落到代码核心是组织好模型前向需要的字段。节点特征x是 [N, d_node]中心性 id 是 [N]空间距离矩阵是 [N, N]边编码输入是 [N, N, path_len, d_edge]再加上注意力掩码attn_bias。批处理时多个图拼成一个大图用attn_bias把不同图之间的注意力遮掉否则图 A 的节点会看到图 B 的节点指标会莫名其妙掉。# 伪代码把预处理字段组织成模型前向所需 batch { x: x, # [N, d_node] in_degree: in_deg, # [N] out_degree: out_deg, # [N] spatial_pos: spatial, # [N, N] edge_input: edge_input, # [N, N, path_len, d_edge] attn_bias: attn_bias, # [N, N] 或 [B, H, N, N] } # 模型内部节点嵌入 中心性嵌入 # 注意力偏置 空间距离查表 边编码映射 # 最终 logits QK^T / sqrt(d) biasattn_bias是最容易被忽略的字段。官方实现里它既承担图间隔离也承担 padding 掩码。如果自己拼 batch一定要把不同图的节点对位置填成负无穷同一图内部填 0 或实际偏置。spatial_pos的 dtype 要是整型查表时不能是浮点否则会报索引类型错误。edge_input的最后一维要和边编码卷积的输入通道对齐改边特征维度时这里要同步改。实际前向时模型会先算节点表示再把中心性嵌入加进去然后每一层注意力都加上空间和边偏置。偏置是逐层共享还是逐层独立不同实现有差异调参时看配置里的num_spatial和边编码层数。我的建议是先把边编码关掉只留空间编码跑一版确认数据管道没问题再开边编码。否则一旦指标不涨你分不清是边特征没对齐还是模型容量不够。3. 跑通官方实现环境、数据与训练入口的落地顺序Graphormer 官方实现的依赖比普通 PyTorch 项目重图库、化学信息库、训练框架叠在一起版本错一个就可能卡在 import。落地顺序应该是先固定 Python 和 PyTorch再装图库最后装项目本身。数据侧先跑通一个小样本再谈全量训练。3.1 用 conda 固定 Python 与 PyTorch 版本我一般用 conda 建独立环境Python 选 3.9原因是图库和训练框架对 3.9 的 wheel 支持最稳。PyTorch 版本按机器的 CUDA 选不要先装 PyTorch 再换 CUDAwheel 不匹配时错误会藏在运行时报错里。装完 PyTorch 再装 PyTorch Geometric最后装项目依赖。如果仓库有锁文件优先按锁文件不要凭感觉升级。# 以仓库锁文件为准版本按机器 CUDA 改 conda create -n graphormer python3.9 -y conda activate graphormer # CUDA_TAG 用 nvidia-smi 右上角版本对应例如 cu118 export CUDA_TAGcu118 pip install torch --index-url https://download.pytorch.org/whl/${CUDA_TAG} # 先装图库和化学信息库再装训练框架 pip install torch-geometric rdkit fairseq # 项目依赖和本地安装 pip install -r requirements.txt pip install -e .CUDA_TAG不是越新越好要和驱动匹配。pip install torch-geometric装的是纯 Python 部分如果用到稀疏算子还要装对应的 wheel否则运行时会提示找不到扩展。fairseq对 PyTorch 版本比较敏感装完先跑一个最小 import 测试不要等训练启动才发现版本冲突。pip install -e .是为了让项目内的模块可导入如果仓库没有setup.py这一步换成把项目根目录加进PYTHONPATH。装完环境后做三件事python -c import torch; print(torch.cuda.is_available())确认 CUDA 可用python -c import torch_geometric; print(torch_geometric.__version__)确认图库能导入跑仓库自带的单元测试或最小样例。不要跳过最小样例很多维度错误在训练脚本启动几分钟后才暴露。3.2 图数据预处理从 SMILES/边表到 Graphormer 输入Graphormer 的输入不是现成的边表而是带中心性、空间距离、边编码路径的结构化样本。分子图可以从 SMILES 转通用图可以从边表转。预处理要做的节点特征、边索引、边类型、入度出度、最短路径距离、边路径序列。节点特征不要直接拿原子序数当连续值最好做 one-hot 或嵌入边类型要映射成连续 id不能是浮点键型。from rdkit import Chem import torch def smiles_to_graph(smiles, max_nodes128): mol Chem.MolFromSmiles(smiles) if mol is None: return None # 只保留重原子氢原子隐式处理减少节点数 mol Chem.RemoveHs(mol) if mol.GetNumAtoms() max_nodes: return None x [] for atom in mol.GetAtoms(): # 实际项目里建议 one-hot 或嵌入不要直接拿原子序数当连续值 x.append([ atom.GetAtomicNum(), atom.GetDegree(), int(atom.GetIsAromatic()), ]) edge_index, edge_type [], [] for bond in mol.GetBonds(): i, j bond.GetBeginAtomIdx(), bond.GetEndAtomIdx() t int(bond.GetBondTypeAsDouble()) edge_index [[i, j], [j, i]] edge_type [t, t] return { x: torch.tensor(x, dtypetorch.float), edge_index: torch.tensor(edge_index, dtypetorch.long).t().contiguous(), edge_type: torch.tensor(edge_type, dtypetorch.long), }max_nodes是硬截断超过的分子直接丢弃或做子图采样。分子图节点数一般不大128 够覆盖多数样本如果任务里有大分子截断会丢信息最好在预处理阶段统计节点数分布再定。RemoveHs去掉显式氢能显著减少节点数但某些任务氢键很关键去氢前要确认任务需求。边索引正反各存一次edge_type也要正反各一份否则无向图会变成有向图。edge_type用int(bond.GetBondTypeAsDouble())只是示意实际要建一个映射表把单键、双键、三键、芳香键映射到 0 到 N-1。映射表一旦定下训练和推理必须共用不能训练时单键是 0、推理时单键是 1。节点特征里的原子序数最好也做嵌入直接当连续值会让模型误以为原子序数 8 和 9 的距离等于 1 和 2 的距离。3.3 启动训练与日志观察先跑小样本再全量训练入口按仓库实际脚本替换常见是项目自带的train.py或训练框架的命令行。第一次不要全量先抽 1000 个样本batch_size设小num_workers设 0排除数据加载的多进程问题。日志重点看四个量训练 loss、学习率、梯度范数、验证集指标。训练 loss 降但验证指标不动通常是过拟合或掩码有问题梯度范数突然爆掉先降学习率或开梯度裁剪。# TRAIN_ENTRY 按仓库 README 替换先小样本跑通 python ${TRAIN_ENTRY} \ --config configs/graphormer_base.yaml \ --max-train-samples 1000 \ --batch-size 8 \ --num-workers 0 \ --max-epoch 2 \ --log-interval 10max-train-samples是调试用的不是正式训练参数batch-size 8是保守值先确认显存能撑住num-workers 0让数据加载在当前进程跑报错栈更清晰。跑通后把num-workers调到 4 或 8观察 GPU 利用率是否上去。log-interval 10是每 10 个 step 打一次日志小样本下够用。小样本跑通的标准不是 loss 降而是前向不报维度错、反向不报 NaN、验证能跑完、checkpoint 能保存。这四件事都过了再换全量数据。全量训练前把max-train-samples去掉batch-size按显存逐步加每次加一倍观察是否 OOM。不要一上来就多卡单卡跑通再上分布式否则多卡报错会把单卡问题盖住。4. 关键参数怎么设batch size、头数、维度与学习率的取舍表Graphormer 的参数不是越多越好。图数据样本数通常比图像文本少模型容量过大容易过拟合注意力是节点对级别的节点数一大显存涨得比参数量还快。调参顺序我一般按先定最大节点数和 batch size再定模型维度和层数最后调学习率和 dropout。4.1 模型维度与层数的取舍d_model是节点表示维度也是注意力头数的基础。中小图任务从 256 试起节点特征维度低、样本少时 128 也够样本多、边特征复杂再上 512。层数从 6 层试起Graphormer 的注意力偏置计算比普通 Transformer 重层数上去显存和耗时都涨。头数 8 是常见起点d_model能被头数整除每个头 32 维或 64 维。边编码维度不要超过d_model的一半否则边编码会压过节点特征。model: d_model: 256 encoder_layers: 6 encoder_attention_heads: 8 dropout: 0.1 attention_dropout: 0.1 num_spatial: 20 num_edge_types: 8 max_degree: 512num_spatial要和预处理里的max_dist一致设 20 表示距离 0 到 20 各一个嵌入超过 20 共享最后一个 id。num_edge_types要和边类型映射表的大小一致映射表有 8 类就写 8写大了浪费写小了越界。max_degree和 2.1 的截断阈值一致。dropout和attention_dropout先都 0.1过拟合再往上加欠拟合先降。层数和头数的组合可以用小规模对比6 层 8 头、6 层 16 头、12 层 8 头各跑一版小样本看验证集和训练时长。我的经验是图数据上 6 层 8 头经常够用12 层带来的提升不如把边特征做干净。头数加到 16 时每个头维度变小注意力更容易学细碎模式但小数据上不一定稳。4.2 训练参数与显存估算Graphormer 的显存大头是注意力矩阵和边编码。节点数 N 的图注意力矩阵是 N×N多头再乘头数层数再乘系数。batch 里最大节点数决定峰值显存不是平均节点数。估算时按batch_size * heads * max_nodes^2 * 4 bytes * 层数系数粗算系数取 3 到 5再加边编码路径的张量。实际以nvidia-smi为准估算只用来决定从哪个 batch size 起试。参数保守起点加大条件风险batch_size8显存利用率低于 70%太大直接 OOMmax_nodes128数据里 95% 样本低于该值截断丢信息grad_accum1想增大有效 batch 但显存不够训练变慢lr1e-4 到 5e-4小样本、浅层模型太大梯度爆炸warmup总 step 的 5%深层模型太少前期不稳export MAX_NODES128 export BATCH_SIZE16 export GRAD_ACC2 # 有效 batch BATCH_SIZE * GRAD_ACC * GPU_NUMMAX_NODES要和预处理截断一致训练时如果还有超过的样本说明预处理漏了。GRAD_ACC是梯度累积用来在显存不够时凑有效 batch但它不减少单步峰值显存只是把多个小 batch 的梯度加起来。学习率按有效 batch 线性缩放要谨慎图数据上我更倾向先固定 1e-4 到 5e-4看前几百 step 的梯度范数再动。4.3 评估与早停别只看训练 loss图级任务常用 MAE、RMSE、AUC 或准确率方向不同早停逻辑要跟着改。MAE 越小越好AUC 越大越好代码里用一个mode变量控制。验证集指标比训练 loss 更可信训练 loss 降但验证指标连续几个 epoch 不降就该停了。保存 best checkpoint不要只保存 last后面推理和微调都用 best。best float(inf) patience 10 bad 0 mode min # MAE 用 minAUC 用 max for epoch in range(max_epoch): train_one_epoch() metric evaluate() improved metric best if mode min else metric best if improved: best metric bad 0 save_checkpoint() else: bad 1 if bad patience: breakpatience设 10 是经验值小数据可以降到 5大数据可以加到 20。mode一定要和指标方向一致MAE 用 min、AUC 用 max写反了会保存最差模型。evaluate()里要确保模型切到 eval 模式dropout 关掉否则验证指标会抖。早停不是越早越好图数据训练前期指标可能平台很久patience 太小会错过后面的下降。5. 避坑与排查从 OOM 到指标不涨的 5 条血泪记录5.1 显存炸了先查边编码和 batch 组成现象训练启动几秒后 CUDA out of memory报错指向注意力或边编码。原因通常不是模型太大而是 batch 里混进了超大图最大节点数把峰值显存拉爆或者边编码路径长度设得太长N×N×path_len的张量比注意力还大。解决先打印每个 batch 的最大节点数按节点数分桶把超大图单独放小 batch再把路径长度从 10 降到 5 试一版。我一般会在数据加载器里加一个max_nodes_in_batch断言超过就拆 batch比事后调 batch size 稳。5.2 loss 正常但验证指标不动检查掩码与 padding现象训练 loss 稳定下降验证集 MAE 或 AUC 几乎不动甚至变差。原因常见是attn_bias没把 padding 位置遮住模型在 padding 节点上也能注意学到的是噪声或者验证集和训练集的预处理不一致边类型映射表不同。解决取一个 batch 手动打印attn_bias确认 padding 位置是负无穷再检查训练和验证是否共用同一个边类型映射和节点特征归一化。掩码问题很隐蔽模型不会报错只会静默学错。5.3 多卡训练速度不升反降数据并行粒度问题现象单卡一个 epoch 20 分钟双卡变 25 分钟GPU 利用率反而低。原因可能是每个 GPU 上的 batch 太小通信开销盖过计算或者数据加载器在每张卡上重复预处理CPU 成瓶颈。解决先把单卡 batch size 翻倍看显存是否撑住撑不住再用梯度累积数据预处理结果缓存到磁盘不要每张卡每次重新算最短路径。多卡不是自动加速图数据小 batch 时尤其明显。5.4 自定义数据集报维度错节点特征与边特征对齐现象换自己的数据后报size mismatch或索引越界堆栈指向嵌入层。原因通常是节点特征维度改了但模型配置里的d_node没改或者边类型 id 从 0 开始还是从 1 开始不一致嵌入表少一行。解决在数据集类里加断言检查x.shape[1] config.d_node、edge_type.max() config.num_edge_types、spatial_pos.max() config.num_spatial。这三个断言能拦住大部分维度错比看堆栈快。5.5 推理结果与训练不一致dropout 与池化方式现象训练时验证指标很好导出推理后同一批数据结果对不上。原因常见是推理时忘了model.eval()dropout 还在随机丢或者图级池化用了 mean 但推理脚本用了 sum再或者 checkpoint 加载的是 last 而不是 best。解决推理脚本固定torch.no_grad()和model.eval()池化方式从配置里读不要手写加载 checkpoint 后打印一个 batch 的预测值和训练脚本的验证输出对比差太多就查预处理。6. 进阶用法把预训练权重迁移到小图任务的冻结与解冻小图任务样本少从零训 Graphormer 容易过拟合。更稳的做法是拿预训练权重初始化先冻结编码器大部分层只训任务头和最后几层再逐步解冻。冻结不是一冻到底先冻 4 层跑几个 epoch验证指标不降再解冻最后两层学习率给编码器低一档、任务头高一档。# 冻结编码器前几层只训任务头和最后两层 for name, param in model.named_parameters(): if name.startswith(encoder.layers.): layer_id int(name.split(.)[2]) param.requires_grad layer_id 4 else: param.requires_grad Truelayer_id 4表示前 4 层冻结从第 5 层开始训。层数索引从 0 开始还是从 1 开始不同实现不一样冻结前先打印几个参数名确认。任务头始终可训否则模型学不到新标签。解冻时不要一次全开每次加两层观察验证集是否过拟合。阶段冻结范围编码器 lr任务头 lr观察指标第 1 段前 4 层01e-3验证集是否降第 2 段前 2 层1e-55e-4训练 loss 是否稳第 3 段全解冻5e-55e-5是否过拟合验证微调是否有效不能只看最终指标。我一般固定随机种子跑三组从零训、冻结微调、全量微调比较验证集均值和方差。如果冻结微调只比从零训高 0.1 个点但方差大很多说明收益不稳不如把预处理做干净。还要看训练曲线冻结阶段验证指标应该先升后平解冻后如果训练 loss 降、验证指标反降就是过拟合往回冻一层。从那以后我每次迁移 Graphormer 权重都先跑一遍小样本冻结验证再决定解冻层数不再一上来全量微调。希望帮到你。本文还有配套的精品资源点击获取
