人工智能深度学习机器学习计算机视觉NLP强化学习【免费下载链接】deepmind-researchThis repository contains implementations and illustrative code to accompany DeepMind publications项目地址https://gitcode.com/gh_mirrors/de/deepmind-research点击查看免费下载导读本文以仓库中 ogb_lsc/README.md 为骨架系统梳理 DeepMind 在 OGB Large-Scale ChallengeOGB-LSC中两个赛道的完整参赛实现面向量子化学属性预测的 PCQM4M-LSCpcq/目录与面向亿级学术图节点分类的 MAG240M-LSCmag/目录。你将读到两套可复现的完整流程——从环境搭建、数据下载与预处理、k-fold 划分、构象特征生成到非常深的 Encode-Process-Decode GNN 训练、EMA 与 BGRL 自监督、以及多模型集成推理并附上 配置 与 模型 的源码级参数解读可直接作为大规模图学习工程的落地参考。背景OGB-LSC 与 DeepMind 的参赛方案OGB-LSCOGB Large-Scale Challenge是 Open Graph Benchmark 社区组织的超大规模图学习竞赛。DeepMind 的参赛代码位于本仓库的ogb_lsc/目录下覆盖两个赛道PCQM4M-LSC量子化学根据分子图预测分子性质HOMO-LUMO 能隙MAE 指标数据规模达数百万分子代码位于 ogb_lsc/pcq/MAG240M-LSC学术图在包含 2.44 亿节点、17.28 亿边的微软学术图上进行论文主题分类代码位于 ogb_lsc/mag/。两者的核心思想在配套技术报告arXiv:2107.09422标题为Large-scale graph representation learning with very deep GNNs and self-supervision中有完整阐述仓库则提供了可运行的实现。代码结构上两个目录各自独立包含config.py实验配置、experiment.pyJaxline 实验入口、model(s).py模型定义、datasets.py数据管道、ensemble_predictions.py集成推理以及一组run_*.sh/download_*.py脚本主 README 即按pcq/与mag/两大模块组织ogb_lsc/README.md。一、PCQM4M-LSC量子化学分子属性预测1.1 任务与最终成绩方式PCQ 团队的最终测试集成绩来自20 个模型10 折 × 2 个随机种子的集成pcq/README.md。官方记录的训练与推理资源为训练单模型 48 小时使用 4× Google Cloud TPUv4 1× AMD EPYC 7B12 64 核 CPU 2.25GHz推理 6 小时使用 1× NVIDIA V100 16GB GPU 1× Intel Xeon Gold 6148 20 核 CPU 2.40GHz。这些数字对应官方提交时的硬件规模仓库提供的脚本为单机示意实现见下文重训部分说明在普通单卡 GPU 上需要相应调小模型与批大小。1.2 环境搭建在克隆的仓库根目录下创建独立虚拟环境并安装依赖pcq/README.mdpython3 -m venv /tmp/pcq_venv source /tmp/pcq_venv/bin/activate pip3 install --upgrade pip setuptools wheel pip3 install -r ogb_lsc/pcq/requirements.txtogb_lsc/pcq/requirements.txt 中除常规依赖外值得注意两点一是通过 git 固定了jaxline、graph_nets、jraph的特定 commit保证与实验代码的 API 兼容二是包含rdkit-pypi2021.3.2.3构象特征生成依赖与jaxlib0.1.67cuda110CUDA 11.0 构建的 JAX 库。1.3 数据下载与预处理k-fold 划分 构象特征训练中使用的两类额外特征都可以本地生成或直接下载# 方式一本地生成 /bin/bash run_preprocessing.sh -r ${HOME}/pcq/ # 方式二直接下载官方已处理好的数据 python download_pcq.py --task_root${HOME}/pcq/ --payloaddatarun_preprocessing.sh内部完成两步ogb_lsc/pcq/run_preprocessing.sh调用generate_validation_splits.py将官方验证集380,670 个样本见 datasets.py均分为 10 份K 10写入${TASK_ROOT}/data/k_fold_splits/{0..9}.pklgenerate_validation_splits.py。每个 fold 训练时保留其中 9 份含官方 train split留出 1 份做早停验证调用generate_conformer_features.py --splitstest valid train --num_parallel_procs32用 RDKit 为每个 SMILES 字符串计算构象conformer三维坐标缓存为${TASK_ROOT}/data/preprocessed/smile_to_conformer.pkl。构象生成是 CPU 密集任务脚本刻意使用multiprocessing.Pool而非多线程以规避 GILgenerate_conformer_features.py。其底层流程在 conformer_utils.py 中Chem.AddHs加氢后用 RDKitAllChem.ETKDGv3距离几何约束嵌入构象随机种子 45、pruneRmsThresh0.01再依次尝试 MMFF 与 UFF 力场做能量最小化最后对齐构象并移除氢原子以匹配图节点数。构象特征以原子 x/y/z 坐标形式进入模型失败样本以 NaN 标记约 0.13%~0.16% 的节点无法生成构象。1.4 模型架构Encode-Process-Decode 与配置逐项解读模型定义在 ogb_lsc/pcq/model.py 的GraphPropertyEncodeProcessDecode类中属于经典的 encode-process-decode 范式Encoder将原子 one-hot 特征173 维、键 one-hot 特征13 维与可选的几何特征分别经 MLP 嵌入到latent_size维Processor堆叠num_message_passing_steps个消息传递层默认 32 层processor_moderesnet表示每层输出与输入做残差相加model.py这是非常深 GNN的关键——残差连接让 32~50 层深的 GNN 可以稳定训练Decoder节点特征按图求和segment_sum后经 MLP 回归出分子级属性同时输出原子与键的 one-hot 辅助预测头用于辅助自监督损失aux_multiplier0.1model.py。默认配置集中在 ogb_lsc/pcq/config.py以下为核心参数及含义参数默认值说明mlp_hidden_size/mlp_layers512 / 2各模块 MLP 的隐藏宽度与层数latent_size512节点/边/全局隐向量维度num_message_passing_steps32消息传递处理器层数processor_moderesnetresnet残差或mlp直连dropedge_rate/dropnode_rate0.1 / 0.1训练时边/节点消息 dropoutadd_relative_distance/add_relative_displacementTrue / True是否把构象几何作为边特征相对位移向量 相对距离标量position_normalization/relative_displacement_normalization2.0 / 1.0绝对/相对坐标归一化系数aux_multiplier0.1原子/键辅助重建损失的权重loss_config_nameRegressionLossConfig回归损失exponent1即 L1MAE目标按NORMALIZE_TARGET_MEAN5.6909、STD1.1561归一化datasets.py训练侧的关键配置优化器为 Adamb1.9, b2.95学习率走optax.warmup_cosine_decay_schedule预热 5e4 步至峰值 1e-4再余弦退火 5e5 步到 0总训练步数 5e6开启EMAema_decay0.9999维护参数滑动平均用于评估config.pysample_random0.05表示训练时随机用 5% 的原子/键特征替换为随机类别一种噪声正则动态批大小按平均每图 16 节点、40 边设定n_node16*64、n_edge40*64、n_graph64config.py。训练评估指标为maebest_model_eval_metric_higher_is_betterFalse越低越好。构象特征在数据管道中还会做增强去均值居中后训练时乘以随机旋转矩阵可附带随机镜像对称使模型对分子的空间朝向具有不变性dataset_utils.py、conformer_utils.py。1.5 训练 40 个模型conformer 与非 conformer 双轨run_training.sh演示了完整训练编排ogb_lsc/pcq/run_training.sh以BASE_SEED42起两个 seed 偏移0、10× 10 个 k-fold 折共 20 个随机种子每折每种子同时训练两种模型合计40 个模型Conformer 模型num_message_passing_steps32latent_size256、mlp_hidden_size1024保留相对距离/位移几何特征run_training.shNon-Conformer 模型num_message_passing_steps50更深关闭add_relative_distance与add_relative_displacementrun_training.sh。每个模型在官方 train 90% 验证集上训练用留出的 10% 验证数据做早停。脚本同时说明了几条实用的中断/续训技巧run_training.shCtrlC仅保存模型不中断训练Ctrl\保存并退出之后可通过--config.restore_path恢复继续训练。experiment.py中的信号处理器实现了这一行为experiment.py。注意官方提交时使用 4× TPUv4 做批并行等效批大小扩大 8 倍、独立的 V100 机器跑在线早停评估训练速度约 4~6 步/秒、50 万步约 1.5 天run_training.sh。仓库脚本仅用于单机示意README 明确提示需要按显存调小批大小或模型规模。1.6 复现最终结果预训练权重 集成推理官方提供了约150 GB的预训练 checkpoint下载与评估命令python download_pcq.py --task_root${HOME}/pcq/ --payloadmodels /bin/bash run_pretrained_eval.sh -r ${HOME}/pcq/run_pretrained_eval.sh的评估逻辑ogb_lsc/pcq/run_pretrained_eval.sh先对test split 在线生成并缓存构象特征generate_conformer_features.py --splitstest脚本刻意不用下载的构象数据关闭 XLA 显存预分配export XLA_PYTHON_CLIENT_PREALLOCATEfalse以让两个模型共享 GPU对 10 折 × 2 seed 逐一以--jaxline_modeeval加载 checkpoint 推理conformer 模型以filter_in_or_out_samples_with_nans_in_conformersout只评估无 NaN 构象的样本non-conformer 模型以in只评估有 NaN 构象的样本——两者按样本索引互补无重叠最后运行ensemble_predictions.py合并两路预测。集成脚本 ensemble_predictions.py 的核心细节merge_complementary_results断言 conformer 与 non-conformer 预测的样本索引无交集后拼接L133-L150集成方式为**取中位数median**而非均值ensemble_test_predictions预测值先 clip 到[0, 20]再集成、集成后再 clip 一次_CLIP_VALUE20.并校验最终值域不超过[0, 40]、无 NaN/Inf 后才调用 OGB 官方的PCQM4MEvaluator.save_test_submission写出y_pred_pcqm4m.npz提交文件。二、MAG240M-LSC亿级学术图节点分类2.1 任务规模与最终成绩方式MAG240M 是微软学术图Microsoft Academic Graph的 2.4 亿节点子集DeepMind 的最终成绩由10 个 k-fold 模型的集成取得mag/README.md。官方记录的资源为训练单模型 72 小时4× Google Cloud TPUv4 1× AMD EPYC 7B12 64 核 CPU 2.25GHz推理 12 小时4× NVIDIA V100 16GB GPU 1× Intel Xeon Gold 6148 20 核 CPU 2.40GHz。数据规模在 data_utils.py 中有硬编码常量可作交叉验证论文节点 121,751,666、作者 122,383,112、机构 25,721合计 244,160,499 节点边总数 1,728,364,232类别数 153。原始节点特征为 768 维 BERT 嵌入EMBEDDING_SIZE 768。2.2 环境搭建与 PCQ 流程一致另需安装与 V100 匹配的 jaxlibmag/README.mdpython3 -m venv /tmp/mag_venv source /tmp/mag_venv/bin/activate pip3 install --upgrade pip setuptools wheel pip3 install -r ogb_lsc/mag/requirements.txt pip install --upgrade jax jaxlib0.1.67cuda110 -f https://storage.googleapis.com/jax-releases/jax_releases.htmlogb_lsc/mag/requirements.txt 相比 PCQ 多出scipy稀疏矩阵、annoy近似最近邻用于邻居构建、graph_nets与jraph的固定 commit。官方 README 还提示如无法通过 pip 找到匹配的 jaxlib 构建可参考 JAX 官方 issue #5231 中的讨论。2.3 数据整理与预处理管线MAG240M 官方数据体积庞大官方流程分为三步第 1 步用竞赛工具包将数据集下载到本地目录ROOT。第 2 步整理为扁平化目录结构/bin/bash organize_data.sh -r ROOT完成后在ROOT/mag240m_kddcup2021/raw下得到node_feat.npy768 维节点特征、node_label.npy、node_year.npy、三类边author_affiliated_with_institution_edges.npy/author_writes_paper_edges.npy/paper_cites_paper_edges.npy以及train_idx.npy/valid_idx.npy/test_idx.npy划分索引即所谓的 raw 数据。需要说明README 中该命令写作-r ROOT而仓库内 organize_data.sh 实际通过-i 下载数据目录与-o 任务根目录两个参数接收路径并以split_and_save_indices.py生成三个划分索引文件。使用前请以脚本内的Usage提示为准。第 3 步运行预处理/bin/bash run_preprocessing.sh -r ROOTrun_preprocessing.sh 依次调用四个构建器csr_builder.py把原始边索引构建为 CSR 格式的稀疏邻接含论文-论文双向边paper_paper_b(_t).npz、作者-机构、作者-论文等见 data_utils.pypca_builder.py从 100 万篇论文的 BERT 特征上估计 PCA_NUMBER_OF_PAPERS_TO_ESTIMATE_PCA_ON 1000000把 768 维压缩为 129 维并由论文特征聚合出作者与机构特征pca_builder.pyneighbor_builder.py用 Annoy 在论文 PCA 特征空间建索引寻找论文的近邻并做标签融合label fusion产出neighbor_indices.npy、fused_node_labels.npy等neighbor_builder.pygenerate_validation_splits.py生成 10 折 k-fold 验证划分。预处理非常耗时耗内存README 明确建议直接下载官方处理好的数据用于训练与评估python3 download_mag.py --task_root${HOME}/mag --payloaddatadownload_mag.py 与 PCQ 版本同源均从deepmind-ogb-lscGCS bucket 匿名下载带 crc32c 校验和与最多 5 次重试download_mag.py。2.4 模型架构InteractionNetwork BGRL 自监督模型定义在 ogb_lsc/mag/models.py 的NodePropertyEncodeProcessDecode类Encoder节点输入为 one-hot 类型论文/作者/机构、one-hot 深度、年份与 PCA 特征等经 MLP 嵌入datasets.pyProcessornum_message_passing_steps默认 4层jraph.InteractionNetwork每层做残差连接graph.nodes residual_graph.nodesmodels.py支持disable_edge_updatesTrue只更新节点表示与use_sent_edgesTrue节点更新时计入发出的消息输出头分类 logits 之外还带projector/predictor两个 MLP用于BGRLBootstrapped Graph Latents自监督对图做两路随机破坏特征以 0.4 概率置零、边以 0.2 概率替换为自环见 losses.py 的get_corrupted_view在线预测器去回归另一路的目标投影损失为对称的 L2 余弦距离symmetrizeTruelosses.py。MAG 配置见 ogb_lsc/mag/config.py要点参数默认值说明mlp_hidden_sizes/latent_size[512] / 256MLP 宽度与隐向量维度num_message_passing_steps4消息传递层数dropout_rate/dropedge_rate0.3 / 0.25节点特征 / 边消息 dropoutnormalization_typelayer_norm支持layer_norm/batch_norm/noneaggregation_functionsum支持sum/mean/softmaxuse_sent_edges/disable_edge_updatesTrue / True消息传递结构开关bgrl_loss_configfeature_drop_prob0.4, edge_drop_prob0.2, symmetrizeTrueBGRL 自监督破坏参数优化器AdamWweight_decay1e-5, b10.9, b20.999base_lr1e-2预热 5e4 步training_steps500,000总训练步数早停指标accuracy越高越好数据管道上MAG 最独特的处理是在线子图采样online_subsampling_kwargsconfig.py为每个中心节点按边类型限定邻居数量max_nb_neighbours_per_type如论文最多 40 个论文邻居 20 个作者邻居并remove_future_nodesTrue剔除发表年份晚于中心节点的邻居、deduplicate_nodesTrue训练时每个 batch 再采样 10 倍于标注节点的无标签节点ratio_unlabeled_data_to_labeled_data10.0配合 BGRL 自监督充分利用大规模无标签数据。动态批大小默认n_node340*256、n_edge720*256、n_graph256config 注释提示 GPU 显存不足时可将 256 降到 48config.py。2.5 训练 20 个模型与早停run_training.sh演示了完整编排ogb_lsc/mag/run_training.shBASE_SEED100两个 seed 组100~109 与 110~119× 10 折 20 个模型。每个模型在官方 train 90% 验证集上训练用各自留出的 10% 验证折做早停训练期num_eval_iterations_to_ensemble5对 5 个评估轮次的输出做集成来选取模型与早停。脚本同样支持CtrlC保存、Ctrl\保存退出并续训。官方提交时的分布式方案4× TPUv4 批并行等效批大小 8 倍、独立 CPU 预先子采样邻居以持续喂饱 TPU、独立 V100 跑在线评估训练约 2 步/秒、50 万步不足 3 天run_training.sh。experiment.py头部还提供了不使用 shell 脚本、直接以模块方式运行的等价命令模板experiment.py例如# 训练 多线程早停 python -m ogb_lsc.mag.experiment \ --configogb_lsc/mag/config.py \ --config.experiment_kwargs.config.dataset_kwargs.data_root${DATA_ROOT} \ --config.experiment_kwargs.config.dataset_kwargs.k_fold_split_id${K_FOLD_INDEX} \ --config.checkpoint_dir${CHECKPOINT_DIR} \ --jaxline_modetrain_eval_multithreaded并可叠加 OOM 调参选项把training.dynamic_batch_size_config的n_node16320、n_edge34560、n_graph48等效平均批大小约 48或临时用dataset_kwargs.use_dummy_adjacenciesTrue加载假邻接矩阵以缩短启动时间。2.6 复现最终结果预训练权重与集成推理下载预训练权重并评估python3 download_mag.py --task_root${HOME}/mag --payloadmodels /bin/bash run_preprocessing.sh -r ${HOME}/mag/澄清README 原文复现步骤写的是run_preprocessing.sh但该脚本只完成数据预处理见上文 2.3从仓库脚本内容看真正执行预训练模型评估并产出提交文件的入口是 run_pretrained_eval.sh它正是与 PCQ 侧run_pretrained_eval.sh对位的复现脚本。run_pretrained_eval.sh的关键点ogb_lsc/mag/run_pretrained_eval.sh使用官方按折挑选出的最优种子BEST_SEEDS(100 111 102 113 104 105 106 107 108 109)README 注释说明两个 seed 组中每折挑验证集准确率更高的那个测试集上EPOCHS_TO_ENSEMBLE50——对每个中心节点采样 50 个不同子图做多轮推理再集成README 提示可降到 5 以大幅缩短耗时验证集性能略有下降但测试集可能持平每模型单轮在测试集约 10~25 分钟每个 k-fold 模型的预测写入独立predictions_dir最后调用ensemble_predictions.py集成。MAG 的集成策略与 PCQ 不同ensemble_predictions.py将各模型 logits 取 softmax 得到概率对多模型概率做平均软投票再叠加多数投票 概率平局裁决use_mode_break_tie_by_meanTrue得到最终类别验证集预测则按折gathered每个验证样本恰好在某一折中被留出正好由对应模型给出。随后同样调用 OGB 官方 evaluator 写出提交文件。三、两个赛道共通的工程方法论对比 pcq/ 与 mag/ 两套实现可以提炼出 DeepMind 大规模 GNN 参赛方案中反复出现的几个工程模式Encode-Process-Decode 残差深堆叠两个模型都采用这一范式PCQ 堆到 32~50 层、MAG 为 4 层但配合大规模子图采样残差连接是深 GNN 可训练的前提k-fold 多 seed 集成PCQ 为 10 折 × 2 seed × 2 模型族conformer/non-conformer 40 模型、中位数集成MAG 为 10 折 × 2 seed 共 20 模型、概率平均集成。两套ensemble_predictions.py都实现了严格的索引一致性校验与提交文件校验值域、NaN 检查保证集成结果可直接用于官方评测EMA 参数平均两处训练均维护ema_paramsPCQema_decay0.9999MAG 采用随步数退火的 EMA 率base_rate0.999用 EMA 参数做评估与预测显著提升稳定性experiment.py动态批大小图规模方差大两套数据管道都用dynamically_batchbatching_utils.py按n_node/n_edge/n_graph三个维度动态组批配合 padding mask 统计利用率自监督辅助PCQ 用原子/键 one-hot 重建辅助损失MAG 用 BGRL 双路图破坏对比损失均在不增加标注成本的前提下提升表示质量——这也正是技术报告标题中 self-supervision 的落地点多阶段交付每套方案都拆成download/preprocess → train → pretrained_eval → ensemble → submission五段脚本数据与模型权重都可独立下载便于逐步复现。四、引用若在论文或工程文档中引用该工作官方建议使用以下 BibTeXogb_lsc/README.mdarticle{deepmind2021ogb, author {Ravichandra Addanki and Peter Battaglia and David Budden and Andreea Deac and Jonathan Godwin and Thomas Keck and Wai Lok Sibon Li and Alvaro Sanchez-Gonzalez and Jacklynn Stott and Shantanu Thakoor and Petar Veli\v{c}kovi\{c}}, title {Large-scale graph representation learning with very deep GNNs and self-supervision}, year {2021}, journal{arXiv preprint arXiv:2107.09422}, }技术细节数据集统计、消融、集成策略的动机可查阅上述技术报告arXiv:2107.09422仓库代码则是与该报告对应的可运行实现。赞分享人工智能深度学习机器学习计算机视觉NLP强化学习【免费下载链接】deepmind-researchThis repository contains implementations and illustrative code to accompany DeepMind publications项目地址https://gitcode.com/gh_mirrors/de/deepmind-research点击查看免费下载相关推荐DGL 实现 OGB-LSCKDD Cup 2021基线MAG240M 节点分类与 PCQM4M 图回归实战指南DGL 实现 OGB LSCKDD Cup 2021基线MAG240M 节点分类与 PCQM4M 图回归实战指南 本指南以 DGL 仓库 examples人工智能机器学习深度学习图计算DGL GraphBolt 异构图表征学习实战用 RGCN 在 ogbn-mag 与 ogb-lsc-mag240m 上做节点分类DGL GraphBolt 异构图表征学习实战用 RGCN 在 ogbn mag 与 ogb lsc mag240m 上做节点分类 本文以 DGL 仓库中 e人工智能机器学习深度学习图计算基于 DGL 的 R-GCN 异构图节点分类从 ogbn-mag 到 ogb-lsc-mag240m 的完整实战基于 DGL 的 R GCN 异构图节点分类从 ogbn mag 到 ogb lsc mag240m 的完整实战 R GCNRelational Graph人工智能机器学习深度学习图计算上一篇gopass 安全模型与安全工程实践威胁模型、加密机制与防篡改设计全解析下一篇终极指南如何优化h2ogpt响应时间目标管理与性能调优 创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考