DGL GraphBolt Layer-Neighbor SamplingLABOR实战用 layer_dependency 与 batch_dependency 化解 GNN 邻域爆炸【免费下载链接】dglPython package built to ease deep learning on graph, on top of existing DL frameworks.项目地址: https://gitcode.com/gh_mirrors/dg/dglLayer-Neighbor Sampling简称 LABOR源自 NeurIPS 2023 论文《Layer-Neighbor Sampling -- Defusing Neighborhood Explosion in GNNs》是 DGL 中针对多层 GNN 小批量训练设计的邻域采样方案。本文以 DGL 仓库中的官方示例 examples/graphbolt/pyg/labor/README.md 为主线结合 GraphBolt 的 LayerNeighborSampler 源码实现与配套测试完整讲解其两个核心参数layer_dependency与batch_dependency的原理、命令行用法、硬件部署模式选择以及如何在ogbn-products等大规模数据集上通过采样节点数与缓存命中率的对比直观感受 LABOR 相对传统 Neighbor Sampling 的性能收益。一、背景多层 GNN 的邻域爆炸以及 LABOR 为何能拆弹GNN 的小批量训练通常按层逐跳采样邻居第 1 层为种子节点采样邻居第 2 层再为这些邻居采样邻居依此类推。随着层数加深参与计算的节点集合呈指数级膨胀即邻域爆炸Neighborhood Explosion问题。经典的 GraphSAGE 式 Neighbor Sampler 在每个批次、每一层都独立做均匀/随机采样导致不同层之间、不同批次之间采到的邻居大量重复冗余算力被浪费在重复计算上。LABORLayer-Neighbor采样通过控制随机数random variates的依赖关系来化解这一问题其思想可归纳为两点跨层依赖layer dependency让各层使用同一组随机数保证同一顶点在每一层采到相同的邻域从而把采样退化为一种子图采样显著减少跨层重复采样的节点数量跨批次依赖batch dependency让相邻 minibatch 使用的随机数缓慢变化论文中记为 κ提高时间局部性temporal locality配合 GPU/CPU 特征缓存可以大幅降低缓存未命中率加速特征向 GPU 的搬运。DGL 官方在仓库中提供了完整可复现的 LABOR 示例位于 examples/graphbolt/pyg/labor使用 PyG 的 GraphSAGE 模型在 GraphBolt 数据管道上完成节点分类训练并公开了采样节点数与缓存未命中率的实测对比数据。二、示例概览目录结构与数据管道示例目录共 4 个文件职责清晰文件作用README.md官方说明参数介绍、性能建议与实测输出node_classification.py主脚本参数解析、GraphSAGE 模型、训练/评估/推理全流程load_dataset.py数据集加载GraphBolt 内置数据集与 DGL 传统数据集cora/citeseer/pubmed/reddit/yelp/flickr统一封装sage_conv.py自定义 SAGEConv 变体--sage-model-variantcustom时使用激活函数为 GELU 并带投影数据管道的组装逻辑集中在 node_classification.py 的create_dataloader中是一条标准 GraphBolt 流水线gb.ItemSampler按batch_size切分训练/验证样本根据args.mode决定是否将数据copy_to目标设备依据--sample-mode调用sample_neighbor传统采样或sample_layer_neighborLABOR 采样后者额外透传layer_dependency与batch_dependency两个参数fetch_feature拉取采样子图的节点特征支持overlap_fetch与 CUDA 流重叠交给gb.DataLoader(num_workers...)驱动整个流水线。三、核心采样器dgl.graphbolt.LayerNeighborSampler示例展示的核心 API 是dgl.graphbolt.LayerNeighborSampler其实现位于 python/dgl/graphbolt/impl/neighbor_sampler.py。从源码结构看它与NeighborSampler同继承自NeighborSamplerImpl唯一的本质区别是底层调用了graph.sample_layer_neighbors定义于 fused_csc_sampling_graph.py而非graph.sample_neighbors。官方注释明确指出LABOR 采用顺序泊松采样sequential Poisson sampling而非普通泊松采样从而让每个顶点被采样的边数保持确定与NeighborSampler一致因此它是NeighborSampler的即插即用替代品drop-in replacement——无需改动模型与数据管道代码只需替换采样器即可。与NeighborSampler相比它在多层 GNN 场景下采样更少的顶点和边且不损害按训练迭代数衡量的收敛速度。其构造函数签名为gb.LayerNeighborSampler( datapipe, graph, fanouts, replaceFalse, prob_nameNone, deduplicateTrue, layer_dependencyFalse, batch_dependency1, overlap_fetchFalse, num_gpu_cached_edges0, gpu_cache_threshold1, cooperativeFalse, asynchronousFalse, )其中layer_dependency与batch_dependency是本例讲解的核心二者共同决定了随机数的生成与复用策略。四、参数一layer_dependency —— 让各层共享同一组随机数启用方式为命令行开关--layer-dependency。启用后采样所用的随机数在各层之间完全相同从而保证同一个顶点在每一层都获得完全相同的邻域。其效果可以从源码中的种子管理逻辑得到印证在NeighborSamplerImpl.sampling_stages中只有当layer_dependencyFalse时每采样完一层才会执行_increment_seed对随机种子 1见 neighbor_sampler.py一旦启用该参数所有层共用同一个随机种子采样结果自然层间一致。这一设计带来两方面的收益采样规模显著下降由于后续层与前面层共享邻域结构跨层的重复节点大幅减少。文档给出的实测是启用--layer-dependency后ogbn-products上每批采样节点数约为190k而未启用时为250k语义上等价于子图采样官方注释称之为 turns LayerNeighborSampler into a subgraph sampling method即后期层被保证采样到与前期层重叠的邻居计算依赖图被压缩。五、参数二batch_dependency —— 让相邻批次共享随机数喂饱缓存--batch-dependencykk 为任意非负整数默认 1实现的是论文《Cooperative Minibatching in Graph Neural Networks》arXiv:2310.12403第 3.2 节提出的依赖式小批量采样dependent minibatching论文中以 κ 表示。它让跨 minibatch 使用的随机数彼此相关从而提升被访问节点和边的时间局部性当配合特征缓存使用时可以观察到一个直接可量化的收益——缓存未命中率cache miss rate随 batch_dependency 增大而显著下降进而加快嵌入特征向 GPU 的传输速度。其底层机制体现在 neighbor_sampler.py 的_init_seed与_set_seed中每个批次携带一份_random_seed与一个插值系数_seed2_contribution随机种子并非每个批次都完全重采而是每 k 个批次才更新一次代码中cnt [-1, batch_dependency]当cnt[0] % cnt[1] 0时把当前种子的最后一个元素滚动到首位再对新位置重新随机self.random_seed[0] self.random_seed[-1]; self.random_seed[-1:].random_(...)批与批之间通过_seed2_contribution (cnt[0] % cnt[1]) / cnt[1]实现随机数的渐进过渡即随机数随时间缓慢漂移而非突变。官方注释对此的概括是设置batch_dependencyκ相当于把随机数变化的速度按 1/κ 比例放缓。更通俗地说相邻 k 个批次看到的采样分布高度相似因此它们访问的节点/边集合大量重叠——这正是缓存所欢迎的访问模式。文档实测数据ogbn-products--num-gpu-cached-features500000batch_dependency缓存未命中率162%3222%仅仅把 k 从 1 提到 32缓存未命中率就从 62% 降到 22%特征读取压力下降近三分之二单轮训练耗时也从约 3.77s 降至约 3.53s。六、命令行参数全表与默认值主脚本 node_classification.py 通过argparse暴露了完整的可调参数整理如下默认值均来自源码参数默认值说明--epochs9999999训练轮数上限配合早停使用--lr0.001Adam 学习率--num-hidden256隐藏层维度--dropout0.5Dropout 概率--batch-size1024训练批次大小--num-workers0数据加载 worker 数--datasetogbn-products可选 ogbn-arxiv / ogbn-products / ogbn-papers100M / igb-hom-tiny~large / reddit / yelp / flickr 等--fanout10,10,10每层采样扇出逗号分隔长度决定 GNN 层数--modepinned-pinned-cuda图存储-特征存储-训练设备组合见下文--layer-dependencyFalse布尔开关启用跨层随机数共享--batch-dependency1非负整数 κ控制跨批次随机数依赖--cpu-feature-cache-policyNoneCPU 特征缓存策略s3-fifo / sieve / lru / clock--num-cpu-cached-features0CPU 特征缓存容量特征条数--num-gpu-cached-features0GPU 特征缓存容量特征条数--early-stopping-patience25早停耐心轮数--sample-modesample_layer_neighborsample_neighbor传统或 sample_layer_neighborLABOR--sage-model-variantcustomcustomGELU 投影精度更高或 originalPyG 原版 SAGEConv--precisionhigh矩阵乘精度torch.set_float32_matmul_precision从 load_dataset.py 可见数据集的加载统一由gb.BuiltinDataset完成OGB / IGB 系列而 reddit / cora / citeseer / pubmed / yelp / flickr 则通过gb.LegacyDataset包装 DGL 传统数据集其中 yelp 为多标签任务使用BCEWithLogitsLoss与 F1 评估官方建议 yelp 上使用--dropout0可获得更好的最终验证与测试精度。七、部署模式六种图/特征/训练设备组合--mode的取值遵循图存储-特征存储-训练设备三段式约定cpu表示 CPU 与 RAMpinned表示 RAM 中的页锁定内存cuda表示 GPU 与显存共 6 种组合用于按显存预算逐级升级模式含义适用场景cpu-cpu-cpu图、特征、训练全在 CPU无 GPU 环境无 CUDA 时脚本自动回退到此模式cpu-cpu-cuda图与特征在内存训练在 GPU图/特征均无法放入显存cpu-pinned-cuda图在内存特征在页锁定内存特征从 host 侧快速搬运pinned-pinned-cuda图与特征均在页锁定内存官方默认模式配合 UVA 异步抓取cuda-pinned-cuda图在显存特征在页锁定内存图可放入显存但特征放不下用 GPU 特征缓存cuda-cuda-cuda图、特征、训练全部在 GPU显存充裕时的最快配置文档给出的性能调优建议按显存从宽到紧依次是--torch-compile脚本内默认对训练/评估步启用torch.compile获得最佳运行时性能显存有富余时用--modecuda-cuda-cuda把整个数据集搬进 GPU放不下全部特征时用--modecuda-pinned-cuda --num-gpu-cached-featuresN图放显存、特征放内存并缓存 N 条节点特征在 GPU图也放不下时用--modepinned-pinned-cuda --num-gpu-cached-featuresN还想把采样本身放到 CPU 执行则用--modecpu-pinned-cuda --num-gpu-cached-featuresN。同时从 node_classification.py 可以看到当--num-gpu-cached-features 0且特征不在cuda设备时脚本会调用gb.gpu_cached_feature为特征套上 GPU 缓存层并通过gpu_cached_feature.miss_rate在训练进度条上实时汇报缓存未命中率若设置了--num-cpu-cached-features且特征为DiskBasedFeature还会调用gb.cpu_cached_feature建立 CPU 侧缓存支持 s3-fifo / sieve / lru / clock 策略并同样汇报 miss rate。示例运行输出中的cache_miss0.619正是这一机制的可观测指标。八、运行示例与实测输出解读以下四个示例均使用ogbn-products默认数据集并统一用--num-gpu-cached-features500000在 GPU 上缓存 50 万条节点嵌入。它们以相同的数据与模型为基准分别演示 LABOR 两个参数的效果以及与原始 GraphSAGE 采样器的对比。1baselinebatch_dependency1缓存未命中率 62%python node_classification.py --num-gpu-cached-features500000 --batch-dependency1Training in pinned-pinned-cuda mode. Loading data... The dataset is already preprocessed. Training: 192it [00:03, 50.95it/s, num_nodes247243, cache_miss0.619] Evaluating: 39it [00:00, 76.01it/s, num_nodes137466, cache_miss0.621] Epoch 00, Loss: 1.1161, Approx. Train: 0.7024, Approx. Val: 0.8612, Time: 3.7688188552856445s2batch_dependency32缓存未命中率降至 22%python node_classification.py --num-gpu-cached-features500000 --batch-dependency32Training in pinned-pinned-cuda mode. Loading data... The dataset is already preprocessed. Training: 192it [00:03, 54.34it/s, num_nodes250479, cache_miss0.221] Evaluating: 39it [00:00, 84.66it/s, num_nodes135142, cache_miss0.226] Epoch 00, Loss: 1.1288, Approx. Train: 0.6993, Approx. Val: 0.8607, Time: 3.5339605808258057s3layer_dependencyTrue采样节点数从 250k 降到 190kpython node_classification.py --num-gpu-cached-features500000 --layer-dependencyTraining in pinned-pinned-cuda mode. Loading data... The dataset is already preprocessed. Training: 192it [00:03, 54.03it/s, num_nodes191259, cache_miss0.626] Evaluating: 39it [00:00, 79.49it/s, num_nodes108720, cache_miss0.627] Epoch 00, Loss: 1.1495, Approx. Train: 0.6932, Approx. Val: 0.8586, Time: 3.5540308952331543s4对照原始 GraphSAGE Neighbor Sampler采样节点 520k是 LABOR 的 2 倍以上python node_classification.py --num-gpu-cached-features500000 --sample-modesample_neighborTraining in pinned-pinned-cuda mode. Loading data... The dataset is already preprocessed. Training: 192it [00:04, 45.60it/s, num_nodes517522, cache_miss0.563] Evaluating: 39it [00:00, 77.53it/s, num_nodes255686, cache_miss0.565] Epoch 00, Loss: 1.1152, Approx. Train: 0.7015, Approx. Val: 0.8652, Time: 4.211000919342041s四组输出横向对比可以得到三条清晰的结论batch_dependency 主攻访存效率κ 从 1 提升到 32GPU 特征缓存未命中率从 62% 降至 22%每批耗时也相应缩短验证了依赖式采样提升时间局部性的设计目标layer_dependency 主攻计算规模启用后每批参与计算的节点数从约 250k 压缩到约 190k约 -24%验证了跨层随机数共享对削减重复计算的作用LABOR 相对传统采样整体占优同样的模型与数据集传统sample_neighbor每批要采样约 520k 节点是 LABOR约 190k~250k的 2 倍以上且每轮耗时更长。由于文档采用早停--early-stopping-patience25保证模型充分收敛后才汇报精度因此最终精度数字具有可比性官方也欢迎社区贡献进一步提升验证/测试精度的方案。九、源码级原理随机种子如何跨层共享、跨批渐变要真正理解 LABOR 两个参数最直接的方式是阅读 python/dgl/graphbolt/impl/neighbor_sampler.py 中NeighborSamplerImpl对种子的三段式管理初始化_init_seedL683-L691为采样器创建独立 RNG 与初始随机种子种子的长度取决于batch_dependencyκ1 时种子为 2 元素数组否则为 1 元素每批设置_set_seedL693-L705批次计数器自增当cnt[0] % cnt[1] 0时滚动更新随机种子否则沿用上一批的种子同时计算插值系数_seed2_contribution把该系数与种子写入 minibatch_random_seed、_seed2_contribution最终经SamplePerLayer传入底层sample_layer_neighbors每层演进_increment_seedL707-L710在sampling_stagesL745-L779中每完成一层采样若layer_dependencyFalse则执行一次种子 1使各层使用不同随机数若为True则所有层共用同一随机数处理完毕后由_delattr_dependency清理临时属性。配套单元测试 tests/python/pytorch/graphbolt/impl/test_neighbor_sampler.pytest_labor_dependent_minibatching对上述行为做了精确验证测试在batch_dependency100的设定下逐批比较相邻批次最内层采样节点的交集——断言相邻批次的采样节点交集大小不小于最内层 fanout从而量化验证相邻批次采样结果高度重叠这一依赖式采样的核心性质同时断言layer_dependencyTrue时第一层采样的输入节点集合与该层采样子图的original_row_node_ids完全相等共享随机数带来的层间一致性而False时输入节点数严格大于该集合层间冗余。十、小结与上手建议LABOR 为多层 GNN 的小批量训练提供了一套用随机数控制采样重复度的优雅方案layer_dependency压缩单批内的层间冗余batch_dependency压缩批间冗余并直接转化为缓存命中率的提升。在实际使用中官方示例的运行路径可以总结为按显存预算选择--mode从pinned-pinned-cuda起步用--num-gpu-cached-features开启 GPU 特征缓存观察训练进度条中的cache_miss逐步调大--batch-dependency如 1 → 32以观察 miss rate 下降与提速叠加--layer-dependency进一步削减采样节点规模若需与传统方案对比基线用--sample-modesample_neighbor即可在同一数据管道上切换回 GraphSAGE 采样器。建议读者在阅读本文后直接进入 examples/graphbolt/pyg/labor 目录运行上述命令并结合 neighbor_sampler.py 与对应测试深入理解两个参数在底层对随机种子的操控方式从而在自有数据集上做出最有依据的调参选择。【免费下载链接】dglPython package built to ease deep learning on graph, on top of existing DL frameworks.项目地址: https://gitcode.com/gh_mirrors/dg/dgl创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 DGL GraphBolt Layer-Neighbor Sampling（LABOR）实战：用 layer_dependency 与 batch_dependency 化解 GNN 邻域爆炸
DGL GraphBolt Layer-Neighbor SamplingLABOR实战用 layer_dependency 与 batch_dependency 化解 GNN 邻域爆炸【免费下载链接】dglPython package built to ease deep learning on graph, on top of existing DL frameworks.项目地址: https://gitcode.com/gh_mirrors/dg/dglLayer-Neighbor Sampling简称 LABOR源自 NeurIPS 2023 论文《Layer-Neighbor Sampling -- Defusing Neighborhood Explosion in GNNs》是 DGL 中针对多层 GNN 小批量训练设计的邻域采样方案。本文以 DGL 仓库中的官方示例 examples/graphbolt/pyg/labor/README.md 为主线结合 GraphBolt 的 LayerNeighborSampler 源码实现与配套测试完整讲解其两个核心参数layer_dependency与batch_dependency的原理、命令行用法、硬件部署模式选择以及如何在ogbn-products等大规模数据集上通过采样节点数与缓存命中率的对比直观感受 LABOR 相对传统 Neighbor Sampling 的性能收益。一、背景多层 GNN 的邻域爆炸以及 LABOR 为何能拆弹GNN 的小批量训练通常按层逐跳采样邻居第 1 层为种子节点采样邻居第 2 层再为这些邻居采样邻居依此类推。随着层数加深参与计算的节点集合呈指数级膨胀即邻域爆炸Neighborhood Explosion问题。经典的 GraphSAGE 式 Neighbor Sampler 在每个批次、每一层都独立做均匀/随机采样导致不同层之间、不同批次之间采到的邻居大量重复冗余算力被浪费在重复计算上。LABORLayer-Neighbor采样通过控制随机数random variates的依赖关系来化解这一问题其思想可归纳为两点跨层依赖layer dependency让各层使用同一组随机数保证同一顶点在每一层采到相同的邻域从而把采样退化为一种子图采样显著减少跨层重复采样的节点数量跨批次依赖batch dependency让相邻 minibatch 使用的随机数缓慢变化论文中记为 κ提高时间局部性temporal locality配合 GPU/CPU 特征缓存可以大幅降低缓存未命中率加速特征向 GPU 的搬运。DGL 官方在仓库中提供了完整可复现的 LABOR 示例位于 examples/graphbolt/pyg/labor使用 PyG 的 GraphSAGE 模型在 GraphBolt 数据管道上完成节点分类训练并公开了采样节点数与缓存未命中率的实测对比数据。二、示例概览目录结构与数据管道示例目录共 4 个文件职责清晰文件作用README.md官方说明参数介绍、性能建议与实测输出node_classification.py主脚本参数解析、GraphSAGE 模型、训练/评估/推理全流程load_dataset.py数据集加载GraphBolt 内置数据集与 DGL 传统数据集cora/citeseer/pubmed/reddit/yelp/flickr统一封装sage_conv.py自定义 SAGEConv 变体--sage-model-variantcustom时使用激活函数为 GELU 并带投影数据管道的组装逻辑集中在 node_classification.py 的create_dataloader中是一条标准 GraphBolt 流水线gb.ItemSampler按batch_size切分训练/验证样本根据args.mode决定是否将数据copy_to目标设备依据--sample-mode调用sample_neighbor传统采样或sample_layer_neighborLABOR 采样后者额外透传layer_dependency与batch_dependency两个参数fetch_feature拉取采样子图的节点特征支持overlap_fetch与 CUDA 流重叠交给gb.DataLoader(num_workers...)驱动整个流水线。三、核心采样器dgl.graphbolt.LayerNeighborSampler示例展示的核心 API 是dgl.graphbolt.LayerNeighborSampler其实现位于 python/dgl/graphbolt/impl/neighbor_sampler.py。从源码结构看它与NeighborSampler同继承自NeighborSamplerImpl唯一的本质区别是底层调用了graph.sample_layer_neighbors定义于 fused_csc_sampling_graph.py而非graph.sample_neighbors。官方注释明确指出LABOR 采用顺序泊松采样sequential Poisson sampling而非普通泊松采样从而让每个顶点被采样的边数保持确定与NeighborSampler一致因此它是NeighborSampler的即插即用替代品drop-in replacement——无需改动模型与数据管道代码只需替换采样器即可。与NeighborSampler相比它在多层 GNN 场景下采样更少的顶点和边且不损害按训练迭代数衡量的收敛速度。其构造函数签名为gb.LayerNeighborSampler( datapipe, graph, fanouts, replaceFalse, prob_nameNone, deduplicateTrue, layer_dependencyFalse, batch_dependency1, overlap_fetchFalse, num_gpu_cached_edges0, gpu_cache_threshold1, cooperativeFalse, asynchronousFalse, )其中layer_dependency与batch_dependency是本例讲解的核心二者共同决定了随机数的生成与复用策略。四、参数一layer_dependency —— 让各层共享同一组随机数启用方式为命令行开关--layer-dependency。启用后采样所用的随机数在各层之间完全相同从而保证同一个顶点在每一层都获得完全相同的邻域。其效果可以从源码中的种子管理逻辑得到印证在NeighborSamplerImpl.sampling_stages中只有当layer_dependencyFalse时每采样完一层才会执行_increment_seed对随机种子 1见 neighbor_sampler.py一旦启用该参数所有层共用同一个随机种子采样结果自然层间一致。这一设计带来两方面的收益采样规模显著下降由于后续层与前面层共享邻域结构跨层的重复节点大幅减少。文档给出的实测是启用--layer-dependency后ogbn-products上每批采样节点数约为190k而未启用时为250k语义上等价于子图采样官方注释称之为 turns LayerNeighborSampler into a subgraph sampling method即后期层被保证采样到与前期层重叠的邻居计算依赖图被压缩。五、参数二batch_dependency —— 让相邻批次共享随机数喂饱缓存--batch-dependencykk 为任意非负整数默认 1实现的是论文《Cooperative Minibatching in Graph Neural Networks》arXiv:2310.12403第 3.2 节提出的依赖式小批量采样dependent minibatching论文中以 κ 表示。它让跨 minibatch 使用的随机数彼此相关从而提升被访问节点和边的时间局部性当配合特征缓存使用时可以观察到一个直接可量化的收益——缓存未命中率cache miss rate随 batch_dependency 增大而显著下降进而加快嵌入特征向 GPU 的传输速度。其底层机制体现在 neighbor_sampler.py 的_init_seed与_set_seed中每个批次携带一份_random_seed与一个插值系数_seed2_contribution随机种子并非每个批次都完全重采而是每 k 个批次才更新一次代码中cnt [-1, batch_dependency]当cnt[0] % cnt[1] 0时把当前种子的最后一个元素滚动到首位再对新位置重新随机self.random_seed[0] self.random_seed[-1]; self.random_seed[-1:].random_(...)批与批之间通过_seed2_contribution (cnt[0] % cnt[1]) / cnt[1]实现随机数的渐进过渡即随机数随时间缓慢漂移而非突变。官方注释对此的概括是设置batch_dependencyκ相当于把随机数变化的速度按 1/κ 比例放缓。更通俗地说相邻 k 个批次看到的采样分布高度相似因此它们访问的节点/边集合大量重叠——这正是缓存所欢迎的访问模式。文档实测数据ogbn-products--num-gpu-cached-features500000batch_dependency缓存未命中率162%3222%仅仅把 k 从 1 提到 32缓存未命中率就从 62% 降到 22%特征读取压力下降近三分之二单轮训练耗时也从约 3.77s 降至约 3.53s。六、命令行参数全表与默认值主脚本 node_classification.py 通过argparse暴露了完整的可调参数整理如下默认值均来自源码参数默认值说明--epochs9999999训练轮数上限配合早停使用--lr0.001Adam 学习率--num-hidden256隐藏层维度--dropout0.5Dropout 概率--batch-size1024训练批次大小--num-workers0数据加载 worker 数--datasetogbn-products可选 ogbn-arxiv / ogbn-products / ogbn-papers100M / igb-hom-tiny~large / reddit / yelp / flickr 等--fanout10,10,10每层采样扇出逗号分隔长度决定 GNN 层数--modepinned-pinned-cuda图存储-特征存储-训练设备组合见下文--layer-dependencyFalse布尔开关启用跨层随机数共享--batch-dependency1非负整数 κ控制跨批次随机数依赖--cpu-feature-cache-policyNoneCPU 特征缓存策略s3-fifo / sieve / lru / clock--num-cpu-cached-features0CPU 特征缓存容量特征条数--num-gpu-cached-features0GPU 特征缓存容量特征条数--early-stopping-patience25早停耐心轮数--sample-modesample_layer_neighborsample_neighbor传统或 sample_layer_neighborLABOR--sage-model-variantcustomcustomGELU 投影精度更高或 originalPyG 原版 SAGEConv--precisionhigh矩阵乘精度torch.set_float32_matmul_precision从 load_dataset.py 可见数据集的加载统一由gb.BuiltinDataset完成OGB / IGB 系列而 reddit / cora / citeseer / pubmed / yelp / flickr 则通过gb.LegacyDataset包装 DGL 传统数据集其中 yelp 为多标签任务使用BCEWithLogitsLoss与 F1 评估官方建议 yelp 上使用--dropout0可获得更好的最终验证与测试精度。七、部署模式六种图/特征/训练设备组合--mode的取值遵循图存储-特征存储-训练设备三段式约定cpu表示 CPU 与 RAMpinned表示 RAM 中的页锁定内存cuda表示 GPU 与显存共 6 种组合用于按显存预算逐级升级模式含义适用场景cpu-cpu-cpu图、特征、训练全在 CPU无 GPU 环境无 CUDA 时脚本自动回退到此模式cpu-cpu-cuda图与特征在内存训练在 GPU图/特征均无法放入显存cpu-pinned-cuda图在内存特征在页锁定内存特征从 host 侧快速搬运pinned-pinned-cuda图与特征均在页锁定内存官方默认模式配合 UVA 异步抓取cuda-pinned-cuda图在显存特征在页锁定内存图可放入显存但特征放不下用 GPU 特征缓存cuda-cuda-cuda图、特征、训练全部在 GPU显存充裕时的最快配置文档给出的性能调优建议按显存从宽到紧依次是--torch-compile脚本内默认对训练/评估步启用torch.compile获得最佳运行时性能显存有富余时用--modecuda-cuda-cuda把整个数据集搬进 GPU放不下全部特征时用--modecuda-pinned-cuda --num-gpu-cached-featuresN图放显存、特征放内存并缓存 N 条节点特征在 GPU图也放不下时用--modepinned-pinned-cuda --num-gpu-cached-featuresN还想把采样本身放到 CPU 执行则用--modecpu-pinned-cuda --num-gpu-cached-featuresN。同时从 node_classification.py 可以看到当--num-gpu-cached-features 0且特征不在cuda设备时脚本会调用gb.gpu_cached_feature为特征套上 GPU 缓存层并通过gpu_cached_feature.miss_rate在训练进度条上实时汇报缓存未命中率若设置了--num-cpu-cached-features且特征为DiskBasedFeature还会调用gb.cpu_cached_feature建立 CPU 侧缓存支持 s3-fifo / sieve / lru / clock 策略并同样汇报 miss rate。示例运行输出中的cache_miss0.619正是这一机制的可观测指标。八、运行示例与实测输出解读以下四个示例均使用ogbn-products默认数据集并统一用--num-gpu-cached-features500000在 GPU 上缓存 50 万条节点嵌入。它们以相同的数据与模型为基准分别演示 LABOR 两个参数的效果以及与原始 GraphSAGE 采样器的对比。1baselinebatch_dependency1缓存未命中率 62%python node_classification.py --num-gpu-cached-features500000 --batch-dependency1Training in pinned-pinned-cuda mode. Loading data... The dataset is already preprocessed. Training: 192it [00:03, 50.95it/s, num_nodes247243, cache_miss0.619] Evaluating: 39it [00:00, 76.01it/s, num_nodes137466, cache_miss0.621] Epoch 00, Loss: 1.1161, Approx. Train: 0.7024, Approx. Val: 0.8612, Time: 3.7688188552856445s2batch_dependency32缓存未命中率降至 22%python node_classification.py --num-gpu-cached-features500000 --batch-dependency32Training in pinned-pinned-cuda mode. Loading data... The dataset is already preprocessed. Training: 192it [00:03, 54.34it/s, num_nodes250479, cache_miss0.221] Evaluating: 39it [00:00, 84.66it/s, num_nodes135142, cache_miss0.226] Epoch 00, Loss: 1.1288, Approx. Train: 0.6993, Approx. Val: 0.8607, Time: 3.5339605808258057s3layer_dependencyTrue采样节点数从 250k 降到 190kpython node_classification.py --num-gpu-cached-features500000 --layer-dependencyTraining in pinned-pinned-cuda mode. Loading data... The dataset is already preprocessed. Training: 192it [00:03, 54.03it/s, num_nodes191259, cache_miss0.626] Evaluating: 39it [00:00, 79.49it/s, num_nodes108720, cache_miss0.627] Epoch 00, Loss: 1.1495, Approx. Train: 0.6932, Approx. Val: 0.8586, Time: 3.5540308952331543s4对照原始 GraphSAGE Neighbor Sampler采样节点 520k是 LABOR 的 2 倍以上python node_classification.py --num-gpu-cached-features500000 --sample-modesample_neighborTraining in pinned-pinned-cuda mode. Loading data... The dataset is already preprocessed. Training: 192it [00:04, 45.60it/s, num_nodes517522, cache_miss0.563] Evaluating: 39it [00:00, 77.53it/s, num_nodes255686, cache_miss0.565] Epoch 00, Loss: 1.1152, Approx. Train: 0.7015, Approx. Val: 0.8652, Time: 4.211000919342041s四组输出横向对比可以得到三条清晰的结论batch_dependency 主攻访存效率κ 从 1 提升到 32GPU 特征缓存未命中率从 62% 降至 22%每批耗时也相应缩短验证了依赖式采样提升时间局部性的设计目标layer_dependency 主攻计算规模启用后每批参与计算的节点数从约 250k 压缩到约 190k约 -24%验证了跨层随机数共享对削减重复计算的作用LABOR 相对传统采样整体占优同样的模型与数据集传统sample_neighbor每批要采样约 520k 节点是 LABOR约 190k~250k的 2 倍以上且每轮耗时更长。由于文档采用早停--early-stopping-patience25保证模型充分收敛后才汇报精度因此最终精度数字具有可比性官方也欢迎社区贡献进一步提升验证/测试精度的方案。九、源码级原理随机种子如何跨层共享、跨批渐变要真正理解 LABOR 两个参数最直接的方式是阅读 python/dgl/graphbolt/impl/neighbor_sampler.py 中NeighborSamplerImpl对种子的三段式管理初始化_init_seedL683-L691为采样器创建独立 RNG 与初始随机种子种子的长度取决于batch_dependencyκ1 时种子为 2 元素数组否则为 1 元素每批设置_set_seedL693-L705批次计数器自增当cnt[0] % cnt[1] 0时滚动更新随机种子否则沿用上一批的种子同时计算插值系数_seed2_contribution把该系数与种子写入 minibatch_random_seed、_seed2_contribution最终经SamplePerLayer传入底层sample_layer_neighbors每层演进_increment_seedL707-L710在sampling_stagesL745-L779中每完成一层采样若layer_dependencyFalse则执行一次种子 1使各层使用不同随机数若为True则所有层共用同一随机数处理完毕后由_delattr_dependency清理临时属性。配套单元测试 tests/python/pytorch/graphbolt/impl/test_neighbor_sampler.pytest_labor_dependent_minibatching对上述行为做了精确验证测试在batch_dependency100的设定下逐批比较相邻批次最内层采样节点的交集——断言相邻批次的采样节点交集大小不小于最内层 fanout从而量化验证相邻批次采样结果高度重叠这一依赖式采样的核心性质同时断言layer_dependencyTrue时第一层采样的输入节点集合与该层采样子图的original_row_node_ids完全相等共享随机数带来的层间一致性而False时输入节点数严格大于该集合层间冗余。十、小结与上手建议LABOR 为多层 GNN 的小批量训练提供了一套用随机数控制采样重复度的优雅方案layer_dependency压缩单批内的层间冗余batch_dependency压缩批间冗余并直接转化为缓存命中率的提升。在实际使用中官方示例的运行路径可以总结为按显存预算选择--mode从pinned-pinned-cuda起步用--num-gpu-cached-features开启 GPU 特征缓存观察训练进度条中的cache_miss逐步调大--batch-dependency如 1 → 32以观察 miss rate 下降与提速叠加--layer-dependency进一步削减采样节点规模若需与传统方案对比基线用--sample-modesample_neighbor即可在同一数据管道上切换回 GraphSAGE 采样器。建议读者在阅读本文后直接进入 examples/graphbolt/pyg/labor 目录运行上述命令并结合 neighbor_sampler.py 与对应测试深入理解两个参数在底层对随机种子的操控方式从而在自有数据集上做出最有依据的调参选择。【免费下载链接】dglPython package built to ease deep learning on graph, on top of existing DL frameworks.项目地址: https://gitcode.com/gh_mirrors/dg/dgl创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考