ARTICLE DETAIL

资讯详情

深耕网站建设与运营推广的一线实战洞察。

基于 DGL 的 Graphormer 图表示学习实战:在 ogbg-molhiv 上微调 PCQM4Mv2 预训练模型

基于 DGL 的 Graphormer 图表示学习实战:在 ogbg-molhiv 上微调 PCQM4Mv2 预训练模型 基于 DGL 的 Graphormer 图表示学习实战在 ogbg-molhiv 上微调 PCQM4Mv2 预训练模型【免费下载链接】dglPython package built to ease deep learning on graph, on top of existing DL frameworks.项目地址: https://gitcode.com/gh_mirrors/dg/dgl导读本文围绕 DGL 官方示例 examples/core/Graphormer/README.md 展开系统讲解如何在 DGL 生态中复现 Graphormer——一种将图结构信息编码进标准 Transformer 的图表示学习模型。通过本文你将掌握 Graphormer 的三种核心结构编码Degree Encoding、Spatial Encoding、Path Encoding的设计原理、ogbg-molhiv 分子图属性预测任务的完整数据预处理流程以及基于 Hugging Faceaccelerate的多卡混合精度微调与评测方法并能在本仓库中直接运行该示例。Graphormer 是什么把图结构注入 TransformerGraphormer 是一种专门为图结构数据设计的 Transformer 模型其核心思想是不修改 Transformer 的骨干架构而是把图的拓扑结构信息编码成注意力机制可以直接消费的偏置项bias与输入特征。其原始论文为Do Transformers Really Perform Bad for Graph Representation?arXiv:2106.05234。具体而言Graphormer 使用了三种结构编码见 examples/core/Graphormer/README.md 的 Introduction 部分编码作用对象进入模型的方式Degree Encoding度编码单个节点的重要性与节点特征相加作为模型输入Spatial Encoding空间编码节点对之间的最短路径距离作为 self-attention 的偏置项Path Encoding路径编码节点对之间最短路径上的边特征作为 self-attention 的偏置项其中度编码 节点特征构成 Transformer 的输入而空间编码 路径编码作为偏置项注入自注意力模块从而让注意力分数能够感知图的连通性。运行环境与依赖原文档要求三个运行依赖accelerate用于简化多 GPU / 混合精度训练Hugging Face 生态transformers提供优化器AdamW与学习率调度器get_polynomial_decay_schedule_with_warmupogb提供 Open Graph Benchmark 数据集的加载器与评测器此外由于示例基于 PyTorch 且使用了 DGL 的内置模块DegreeEncoder、PathEncoder、SpatialEncoder、GraphormerLayer运行环境还需要安装torch与dgl本身。数据集与任务设定ogbg-molhiv 图属性预测该示例的任务是图属性预测Graph Property Prediction使用的数据集为 OGB 的ogbg-molhiv其统计信息如下源自原文档 Dataset 表格Dataset#Graphs#Node Feats#Edge FeatsMetricogbg-molhiv41,12793ROC-AUC#Graphs41,127分子图总数#Node Feats9每个原子的 9 维特征原子类型等#Edge Feats3每条化学键的 3 维特征键型等MetricROC-AUC评估指标为 ROC-AUC任务本质是二分类分子是否有 HIV 抑制活性。在 examples/core/Graphormer/dataset.py 中MolHIVDataset通过DglGraphPropPredDataset(nameogbg-molhiv)加载数据集并调用dataset.get_idx_split()获取官方划分的 train / valid / test 三个子集。数据预处理从 DGL 图到 Graphormer 输入张量Graphormer 的输入与常规 GNN 不同它要求把所有图 padding 到相同节点数并额外提供距离矩阵、路径信息与注意力掩码。MolHIVDataset的预处理与collate函数examples/core/Graphormer/dataset.py实现了这一转换具体包括1. 最短路径距离与路径计算from dgl import shortest_dist for g, label in dataset: spd, path shortest_dist(g, rootNone, return_pathsTrue) g.ndata[spd] spd g.ndata[path] path这里使用了 DGL 的shortest_dist接口实现在 python/dgl/transforms/functional.py一次性求出spd所有节点对之间的最短路径距离矩阵path对应最短路径上经过的边 ID 序列return_pathsTrue时返回。这两者分别喂给空间编码器与路径编码器。注意shortest_dist对不可达节点对返回-1这一约定贯穿后续 padding 逻辑。2. 填充与掩码构造collate 的核心逻辑# 虚拟节点 1Graphormer 为每个图添加一个虚拟节点 attn_mask th.zeros(num_graphs, max_num_nodes 1, max_num_nodes 1) ... attn_mask[i, :, num_nodes[i] 1:] 1 # 无效位置标记为 True dist -th.ones((num_graphs, max_num_nodes, max_num_nodes), dtypeth.long)关键细节虚拟节点virtual nodeGraphormer 在输入序列最前面追加一个与所有节点相连的虚拟节点其表示作为整图的图嵌入graph embedding输出注意力掩码padding 出的不存在的节点位置在attn_mask中被标记为True避免注意力在这些位置产生计算距离矩阵 padding由于shortest_dist对不可达对返回-1padding 节点对真实节点的距离同样用-1填充度特征截断th.clamp(g.in_degrees() 1, min0, max512)将度截断到[0, 512]区间以匹配DegreeEncoder的表大小路径长度截断每条最短路径只保留前max_len5条边与模型multi_hop_max_dist5一致超出的部分截断不足的部分用-1填充边 ID边特征对齐edata th.cat((edata, th.zeros(1, edata.shape[1])), dim0)在边特征表末尾追加一行全零作为不存在边的占位之后用edata[shortest_path]完成路径边特征的查表。最终collate返回 7 个张量顺序为labels, attn_mask, node_feat, in_degree, out_degree, path_data, dist与训练循环中的解包顺序一一对应见 examples/core/Graphormer/main.py 的train_epoch。模型结构DGL 内置 Graphormer 组件的组装示例模型定义在 examples/core/Graphormer/model.py 中。与从零实现不同该示例直接复用了 DGL 内置的四个图 Transformer 模块from dgl.nn import DegreeEncoder, GraphormerLayer, PathEncoder, SpatialEncoder模型超参数一览参数默认值含义num_classes1二分类输出维度edge_dim3边特征维度对应 ogbg-molhiv 的 #Edge Featsnum_atoms4608原子节点类型的 embedding 表大小max_degree512度编码的上界num_spatial511空间距离编码上界multi_hop_max_dist5路径编码的最大跳数num_encoder_layers12Transformer 编码器层数embedding_dim768特征维度ffn_embedding_dim768FFN 隐藏层维度num_attention_heads32注意力头数dropout0.1Dropout 概率pre_layernormTrue是否 Pre-LN先 LayerNorm 再做注意力/FFNactivation_fnnn.GELU()FFN 激活函数DegreeEncoder度编码DegreeEncoder实现于 python/dgl/nn/pytorch/gt/degree_encoder.py是一个可学习的度嵌入模块directionboth时维护两张nn.Embedding(max_degree 1, embedding_dim)表分别编码入度与出度并相加输入形状为(2, B, N)stack 的入度与出度输出形状为(B, N, d)。在模型前向中deg_emb self.degree_encoder(th.stack((in_degree, out_degree)))后与节点特征相加node_feat self.atom_encoder(node_feat.int()).sum(dim-2) node_feat node_feat deg_emb注意atom_encoder是一个nn.Embedding(num_atoms 1, embedding_dim, padding_idx0)输入为原子类型 ID输出在原子维度上求和把每个节点的多热原子类型向量折叠为单个向量。SpatialEncoder 与 PathEncoder注意力偏置在 examples/core/Graphormer/model.py 的forward中attn_bias th.zeros(num_graphs, max_num_nodes 1, max_num_nodes 1, self.num_heads, devicedist.device) path_encoding self.path_encoder(dist, path_data) spatial_encoding self.spatial_encoder(dist) attn_bias[:, 1:, 1:, :] path_encoding spatial_encodingSpatialEncoderpython/dgl/nn/pytorch/gt/spatial_encoder.py把最短路径距离矩阵查表得到偏置表大小为max_dist 2padding_idx0用于屏蔽距离为 -1 的节点对输入(B, N, N)输出(B, N, N, H)PathEncoderpython/dgl/nn/pytorch/gt/path_encoder.py先对距离做clamp(dist, min1, maxmax_len)再通过einsum(bxyld,lhd-bxyh, path_data, edge_embedding)把每条最短路径上的边特征聚合为路径偏置并除以最短距离做归一化输出同样为(B, N, N, H)。虚拟节点与所有节点的空间偏置通过可学习的graph_token_virtual_distance广播填充到attn_bias的第 0 行与第 0 列见 examples/core/Graphormer/model.py。GraphormerLayer带偏置的多头注意力编码层模型主体由 12 个GraphormerLayer堆叠而成。GraphormerLayer实现于 python/dgl/nn/pytorch/gt/graphormer.py是一个带密集多头注意力BiasedMHA的 Transformer 编码层其关键参数feat_size特征维度hidden_sizeFFN 隐藏维度num_heads注意力头数要求能整除feat_sizeattn_bias_type偏置类型可选add加性偏置或mul乘性偏置默认addnorm_firstPre-LN 与 Post-LN 的选择dropout/attn_dropout残差路径与注意力权重的 DropoutactivationFFN 激活函数。forward的输入为nfeat (B, N, feat_size)、attn_bias (B, N, N, H)与可选的attn_mask (B, N, N)实现标准残差 LayerNorm FFN 结构python/dgl/nn/pytorch/gt/graphormer.py。图表示输出头编码器输出的第 0 个位置虚拟节点即图表示graph_rep x[:, 0, :] graph_rep self.layer_norm(self.activation_fn(self.lm_head_transform_weight(graph_rep))) graph_rep self.embed_out(graph_rep) self.lm_output_learned_bias微调时需调用reset_output_layer_parameters()将输出层embed_out与lm_output_learned_bias重新初始化以适配新的下游任务。运行方式多卡混合精度微调原文档给出的运行命令为accelerate launch --multi_gpu --mixed_precisionfp16 main.py该命令使用 Hugging Faceaccelerate启动训练其中--multi_gpu启用多 GPU 分布式训练--mixed_precisionfp16启用 FP16 混合精度以降低显存占用并加速迭代。脚本还支持的 CLI 参数main.py的argparse还暴露了两个可调参数examples/core/Graphormer/main.py参数默认值说明--seed1随机种子用于复现数据加载顺序--batch_size16每张卡上的 batch size注意原文档 NOTE脚本会自动下载在PCQM4Mv2上预训练好的权重graphormer_pcqm.pth。为了复现原文档报告的结果需要将总 batch size 设置为 64例如单卡batch_size64或多卡每卡 16、共 4 卡。完整训练流程train_val_pipeline结合 examples/core/Graphormer/main.py 的train_val_pipeline整个流程为加载数据集实例化MolHIVDataset()输出 train / val / test 规模构建 DataLoader使用 DGL 的GraphDataLoadercollate_fndataset.collatenum_workers16、pin_memoryTrue下载预训练权重download(urlhttps://data.dgl.ai/pre_trained/graphormer_pcqm.pth)加载后load_state_dict重置输出层model.reset_output_layer_parameters()设定训练超参num_epochs16总更新步数total_updates 33000 * num_epochs / batch_sizewarmup 比例为论文同款0.16优化器与调度器AdamW(lr1e-4, eps1e-8, weight_decay0)get_polynomial_decay_schedule_with_warmuplr_end1e-9, power1.0accelerate.prepare将model, optimizer, train_loader, val_loader, test_loader, lr_scheduler全部交给accelerator.prepare()以适配分布式环境训练与评测每轮执行train_epoch训练集与evaluate_network验证集、测试集记录并打印 ROC-AUC选取最佳轮次以验证集 AUC 最高的轮次对应的测试集结果作为最终报告值。训练循环train_epoch的关键实现细节损失函数为nn.BCEWithLogitsLoss()二分类 logits梯度回传使用accelerator.backward(loss)由accelerate统一处理多卡与混合精度评测使用accelerator.gather_for_metrics聚合所有进程的预测与标签再交给 OGB 的Evaluator(nameogbg-molhiv)计算 ROC-AUC每批结束后手动del各张量并调用th.cuda.empty_cache()释放显存。实验结果与复现说明原文档的 Summary 给出的参考结果为ogbg-molhiv在 PCQM4Mv2 上预训练ROC-AUC ≈0.791复现该结果的前提包括自动下载 PCQM4Mv2 预训练权重、总 batch size 保持为 64、按脚本默认的 16 个 epoch 与 0.16 的 warmup 比例训练。该数值来自示例作者在默认配置下的实测记录实际运行会因硬件、随机种子--seed与框架版本略有浮动。源码验证与延伸阅读该示例的四个 DGL 内置模块均有对应的单元测试覆盖例如 tests/python/pytorch/nn/test_nn.py 中的test_degree_encoder覆盖max_degree、embedding_dim、direction组合与test_GraphormerLayer覆盖attn_bias_type与norm_first组合可用于验证模块的正确性与理解参数行为。如果你想进一步深入学习建议按以下路径在本仓库中展开示例骨架examples/core/Graphormer/main.py、examples/core/Graphormer/model.py、examples/core/Graphormer/dataset.py底层模块实现python/dgl/nn/pytorch/gt/graphormer.py、python/dgl/nn/pytorch/gt/degree_encoder.py、python/dgl/nn/pytorch/gt/spatial_encoder.py、python/dgl/nn/pytorch/gt/path_encoder.py最短路径计算 APIpython/dgl/transforms/functional.py 中的shortest_dist数据加载器python/dgl/dataloading 目录下的GraphDataLoader单卡 / CPU 版本 Graphormer 的类似使用可参考 examples/sparse 下的图 Transformer 示例graph_transformer.py了解同一思想在 DGL Sparse 接口下的另一种实现。总结本文完整覆盖了 examples/core/Graphormer/README.md 的全部内容并深入到底层实现Graphormer 通过度编码增强节点输入、通过空间编码与路径编码把结构信息注入注意力偏置从而在几乎不修改 Transformer 骨架的前提下获得强大的图表示能力。借助 DGL 内置的DegreeEncoder、PathEncoder、SpatialEncoder与GraphormerLayer配合accelerate的多卡混合精度能力你可以在 ogbg-molhiv 上以约 0.791 的 ROC-AUC 复现预训练微调基线并将这套结构编码 偏置注意力的范式迁移到其他图属性预测任务中。【免费下载链接】dglPython package built to ease deep learning on graph, on top of existing DL frameworks.项目地址: https://gitcode.com/gh_mirrors/dg/dgl创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表