简介：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_degree=512): # edge_index: [2, E]，有向边；无向图正反各存一次 src, dst = edge_index[0], edge_index[1] in_deg = torch.bincount(dst, minlength=num_nodes) out_deg = torch.bincount(src, minlength=num_nodes) # 长尾截断：超过 max_degree 的节点共享最大 id in_deg = torch.clamp(in_deg, max=max_degree) out_deg = torch.clamp(out_deg, max=max_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_dist=20): # adj: list[list[int]]，无权无向邻接表 spatial = torch.full((num_nodes, num_nodes), max_dist, dtype=torch.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 spatial

bfs_spatial_pos输出的是节点对最短路径距离矩阵，max_dist既是截断阈值，也是空间嵌入表的行数。超过阈值的节点对统一落到最大 id，不再区分 21 跳还是 30 跳。这个函数复杂度是 O(N*(N+E))，节点数上千时就会明显拖慢预处理。我的经验是：节点数超过 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) + bias

attn_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 再换 CUDA，wheel 不匹配时错误会藏在运行时报错里。装完 PyTorch 再装 PyTorch Geometric，最后装项目依赖。如果仓库有锁文件，优先按锁文件，不要凭感觉升级。

# 以仓库锁文件为准；版本按机器 CUDA 改 conda create -n graphormer python=3.9 -y conda activate graphormer # CUDA_TAG 用 nvidia-smi 右上角版本对应，例如 cu118 export CUDA_TAG=cu118 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_nodes=128): 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, dtype=torch.float), "edge_index": torch.tensor(edge_index, dtype=torch.long).t().contiguous(), "edge_type": torch.tensor(edge_type, dtype=torch.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 10

max-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: 512

num_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 起试。

export MAX_NODES=128 export BATCH_SIZE=16 export GRAD_ACC=2 # 有效 batch = BATCH_SIZE * GRAD_ACC * GPU_NUM

MAX_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 用 min，AUC 用 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: break

patience设 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 = True

layer_id >= 4表示前 4 层冻结，从第 5 层开始训。层数索引从 0 开始还是从 1 开始，不同实现不一样，冻结前先打印几个参数名确认。任务头始终可训，否则模型学不到新标签。解冻时不要一次全开，每次加两层，观察验证集是否过拟合。

验证微调是否有效，不能只看最终指标。我一般固定随机种子跑三组：从零训、冻结微调、全量微调，比较验证集均值和方差。如果冻结微调只比从零训高 0.1 个点，但方差大很多，说明收益不稳，不如把预处理做干净。还要看训练曲线：冻结阶段验证指标应该先升后平，解冻后如果训练 loss 降、验证指标反降，就是过拟合，往回冻一层。

从那以后我每次迁移 Graphormer 权重，都先跑一遍小样本冻结验证，再决定解冻层数，不再一上来全量微调。希望帮到你。

本文还有配套的精品资源，点击获取