ARTICLE DETAIL

资讯详情

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

DGL GraphBolt 快速入门:用数据管道(DataPipe)搭建 GNN 训练 Dataloader

DGL GraphBolt 快速入门:用数据管道(DataPipe)搭建 GNN 训练 Dataloader 人工智能机器学习深度学习图计算【免费下载链接】dglPython package built to ease deep learning on graph, on top of existing DL frameworks.项目地址https://gitcode.com/gh_mirrors/dg/dgl点击查看免费下载GraphBolt 是 DGL 中面向大规模图训练的数据加载解决方案本文以 examples/graphbolt/quickstart 下的两个完整示例为主线讲解如何用dgl.graphbolt以声明式数据管道的方式搭建 GNN 训练所需的 Dataloader一条流水线同时完成小批量切分、设备搬运、邻居采样、负采样与特征抓取。读完本文你将掌握ItemSampler → copy_to → sample_neighbor → sample_uniform_negative → fetch_feature → DataLoader的完整组装方法并能独立复现 2 层 GCN 节点分类与 2 层 GraphSAGE 链接预测两种经典训练范式。引言GraphBolt 解决什么问题GraphBolt 的设计目标非常聚焦——提供创建 Dataloader 以训练图神经网络所需的全部组件原文Graphbolt provides all you need to create a dataloader to train a Graph Neural Networks。它并不重新发明模型层而是把 GNN 训练中最耗时的数据侧流程小批量生成、邻居采样、特征抓取、负采样、数据搬运抽象为一条可组合的数据管道DataPipe让用户以链式调用的方式声明式地组装完整的数据加载链路而不是像传统写法那样手写采样循环、手动管理 batch 与特征索引。本教程配套两个可运行的示例均位于仓库 examples/graphbolt/quickstart 目录node_classification.py在 Cora 数据集上训练一个 2 层图卷积网络GCN完成节点分类link_prediction.py在 Cora 数据集上训练一个 2 层 GraphSAGE 完成链接预测并用 AUROC 评估。两个示例共用同一条建模思路先加载数据集再用graphbolt的 DataPipe 拼装 Dataloader最后迭代 Dataloader 完成训练与评估。下面逐一拆解。运行准备示例依赖dgl.graphbolt随 DGL 一并提供与 PyTorch另外需要torchmetrics节点分类示例的准确率计算和torcheval链接预测示例的 AUROC 计算。进入示例目录后直接运行即可# 节点分类 python examples/graphbolt/quickstart/node_classification.py # 链接预测 python examples/graphbolt/quickstart/link_prediction.py两个脚本都会自动通过gb.BuiltinDataset(cora).load()下载并加载 Cora 数据集无需手工准备数据文件。脚本会自动检测硬件优先使用cuda:0否则回退到 CPUtorch.device(cuda:0 if torch.cuda.is_available() else cpu)。节点分类快速入门2 层 GCN第一步加载数据集import dgl.graphbolt as gb dataset gb.BuiltinDataset(cora).load()gb.BuiltinDataset是OnDiskDataset的子类定义见 python/dgl/graphbolt/impl/ondisk_dataset.py负责从 AWS S3 下载内置数据集并加载为OnDiskDataset。加载完成后数据集对象提供三个关键成员dataset.graphFusedCSCSamplingGraph用于采样的图结构dataset.feature特征存储可按节点/边类型与特征名查询特征dataset.tasks任务列表Cora 数据集包含两个任务——tasks[0]是节点分类tasks[1]是链接预测。从dataset.tasks[0].metadata[num_classes]可以拿到类别数作为模型输出维度输入维度则由特征形状推断in_size dataset.feature.size(node, None, feat)[0] out_size dataset.tasks[0].metadata[num_classes]第二步组装 Dataloader 数据管道这是 GraphBolt 的核心。create_dataloader函数以链式调用把四类 DataPipe 串成一条流水线def create_dataloader(dataset, itemset, device): # 1. 从 itemset 中采样种子节点切成 batch datapipe gb.ItemSampler(itemset, batch_size16) # 2. 将 mini-batch 搬运到指定设备供后续采样与训练使用 datapipe datapipe.copy_to(device) # 3. 为种子节点采样邻居两层fanout 分别为 4 和 2 datapipe datapipe.sample_neighbor(dataset.graph, fanouts[4, 2]) # 4. 抓取采样到的节点的特征 datapipe datapipe.fetch_feature( dataset.feature, node_feature_keys[feat] ) # 5. 实例化为 DataLoader return gb.DataLoader(datapipe)各阶段职责如下gb.ItemSampler(itemset, batch_size16)ItemSet定义了要遍历的样本是什么这里是训练/验证/测试集的种子节点ItemSampler负责将其切成指定大小的小批量。它支持shuffle、drop_last、seed可复现的随机打乱种子等参数详见 python/dgl/graphbolt/item_sampler.py。.copy_to(device)把 mini-batch 搬运到采样与训练所在的设备。注意当copy_to放在管道前部时DataLoader的num_workers必须为 0DataLoader文档明确说明多进程下不支持 CUDA 使用见 python/dgl/graphbolt/dataloader.py。.sample_neighbor(dataset.graph, fanouts[4, 2])邻居采样fanouts的长度即采样层数[4, 2]表示第 1 层为每个节点采样 4 个邻居、第 2 层采样 2 个。注意 fanout 顺序是从最外层到最内层源码注释The fanout order is from the outermost layer to innermost layer。采样输出是紧凑化compacted后的子图每个 batch 对应一层一个SampledSubgraph。采样还可配置replace是否有放回、prob_name按边权重采样、deduplicate跨层种子去重等参数见 python/dgl/graphbolt/impl/neighbor_sampler.py。.fetch_feature(dataset.feature, node_feature_keys[feat])把采样到的节点对应的feat特征取出来装入 mini-batch。gb.DataLoader(datapipe)把整条管道包装为可迭代的数据加载器支持num_workers多进程、persistent_workers、max_uva_threads等参数见 python/dgl/graphbolt/dataloader.py。第三步定义并训练 2 层 GCN模型使用dgl.nn.GraphConv堆叠两层中间加 ReLUclass GCN(nn.Module): def __init__(self, in_size, out_size, hidden_size16): super().__init__() self.layers nn.ModuleList() self.layers.append(dglnn.GraphConv(in_size, hidden_size)) self.layers.append(dglnn.GraphConv(hidden_size, out_size)) def forward(self, blocks, x): hidden_x x for layer_idx, (layer, block) in enumerate(zip(self.layers, blocks)): hidden_x layer(block, hidden_x) is_last_layer layer_idx len(self.layers) - 1 if not is_last_layer: hidden_x F.relu(hidden_x) return hidden_x注意forward的输入blocks它是 mini-batch 中按层组织的采样子图列表与fanouts[4, 2]对应共 2 个 block每一层卷积作用在对应的 block 上——这正是 GraphBolt Dataloader 直接产出的训练就绪格式。训练循环极其简洁迭代 Dataloader 即可拿到模型所需的全部张量for step, data in enumerate(dataloader): x data.node_features[feat] # 采样节点的特征 y data.labels # 种子节点的真实标签 y_hat model(data.blocks, x) # 前向 loss F.cross_entropy(y_hat, y) # 损失 optimizer.zero_grad() loss.backward() optimizer.step()每个 epoch 结束后分别用task.validation_set与task.test_set构建验证/测试 Dataloader 评估准确率torchmetrics.functional.accuracytaskmulticlass。默认训练 10 个 epoch学习率1e-2输出格式为Epoch {epoch:03d} | Loss ... | Val Acc ... | Test Acc ...。链接预测快速入门2 层 GraphSAGE数据管道训练与测试的差异链接预测的数据管道与节点分类类似但多了负采样与种子边排除两个环节且训练/测试时管道不同def create_dataloader(dataset, device, is_trainTrue): task dataset.tasks[1] # 链接预测任务 itemset task.train_set if is_train else task.test_set datapipe gb.ItemSampler(itemset, batch_size256) datapipe datapipe.copy_to(device) if is_train: # 训练为每条种子边采样 1 条负边 datapipe datapipe.sample_uniform_negative( dataset.graph, negative_ratio1 ) # 训练两层邻居采样 datapipe datapipe.sample_neighbor(dataset.graph, fanouts[4, 2]) # 训练从子图中剔除种子边防止标签泄露 datapipe datapipe.transform(gb.exclude_seed_edges) else: # 测试全量邻居fanout-1 表示采样所有邻居 datapipe datapipe.sample_neighbor(dataset.graph, fanouts[-1, -1]) datapipe datapipe.fetch_feature( dataset.feature, node_feature_keys[feat] ) return gb.DataLoader(datapipe)三个训练专属环节的作用sample_uniform_negative(dataset.graph, negative_ratio1)对每条种子正边均匀采样negative_ratio条负边。UniformNegativeSampler的具体实现见 python/dgl/graphbolt/impl/uniform_negative_sampler.py。sample_neighbor(dataset.graph, fanouts[4, 2])与节点分类一致的两层邻居采样。链接预测中采样器会先从正负节点对中收集去重后的节点作为种子再采样NeighborSampler文档对此有专门说明。.transform(gb.exclude_seed_edges)把种子边从采样子图中剔除。exclude_seed_edges定义于 python/dgl/graphbolt/external_utils.py它防止模型在消息传递阶段看见用于训练的正样本边避免信息泄露。测试时则用fanouts[-1, -1]全量邻居采样不做负采样、不排除种子边以保证评估的完整性。模型与损失模型是两层SAGEConv(..., mean)编码器加一个两层 MLP 打分器predictor对源/目标节点嵌入做逐元素乘积后输出单个 logitclass GraphSAGE(nn.Module): def __init__(self, in_size, hidden_size16): super().__init__() self.layers nn.ModuleList() self.layers.append(SAGEConv(in_size, hidden_size, mean)) self.layers.append(SAGEConv(hidden_size, hidden_size, mean)) self.predictor nn.Sequential( nn.Linear(hidden_size, hidden_size), nn.ReLU(), nn.Linear(hidden_size, 1), )训练循环中通过data.compacted_seeds.T拿到紧凑化后的正负节点对索引进而取出两端嵌入compacted_seeds data.compacted_seeds.T y model(data.blocks, x) logits model.predictor( y[compacted_seeds[0].long()] * y[compacted_seeds[1].long()] ).squeeze() loss F.binary_cross_entropy_with_logits(logits, labels.float())评估测试阶段用torcheval.metrics.BinaryAUROC计算 AUROC把测试 Dataloader 产出的所有 logit 与标签拼接后一次性更新指标并compute()。注意示例文件头部有免责声明——该示例没有从原始图中剔除测试边可能存在数据泄露作者明确表示We are ignoring this issue for this example because we are focused on demonstrating usability即仅用于演示 GraphBolt 的易用性正式实验需自行处理边划分。深入理解数据管道背后的核心对象MiniBatch贯穿全流程的统一数据结构MiniBatch定义见 python/dgl/graphbolt/minibatch.py是数据加载过程中所有阶段的输入输出统一结构上面示例中反复使用的字段包括字段含义seeds种子项一维张量表示种子节点二维张量每行表示一条边/超链接等labels与种子对应的标签节点分类时是节点类别链接预测时是边标签sampled_subgraphs每层采样一个SampledSubgraph按层组织input_nodes最外层所有采样层涉及的全部输入节点node_features抓取到的节点特征字典键为特征名异构图下为(节点类型, 特征名)元组compacted_seeds紧凑化后的种子对应采样子图中的新索引blocks由sampled_subgraphs转换得到的 DGL block 列表可直接喂给 GNN 层节点分类中只用到node_features、labels、blocks链接预测额外用到compacted_seeds做节点对索引。ItemSet 与 ItemSamplergb.ItemSet见 python/dgl/graphbolt/itemset.py是样本集合的轻量包装支持四种形态单个整数等价于torch.arange生成的节点序列单个张量按第一维索引的节点集合张量元组同形状如(节点, 标签)配对张量元组不同形状如(边, 标签)配对逐项对应切分。每个 ItemSet 可以指定names如seeds、labels这些名字与MiniBatch的属性名对齐从而让ItemSampler自动把切好的 batch 组装成带语义的MiniBatch。ItemSampler本身基于torch.utils.data.IterDataPipe实现因此可以与 PyTorch 官方的任意迭代型 DataPipe 继续拼接见 python/dgl/graphbolt/item_sampler.py 的说明。DataLoader 的多进程策略gb.DataLoader见 python/dgl/graphbolt/dataloader.py会在num_workers 0时自动改造数据管道在ItemSampler后插入sharding_filter()实现各 worker 均匀分配 mini-batch并在FeatureFetcher处切断管道——特征抓取之前的阶段采样等在子进程中执行特征抓取之后的阶段回到主进程。这个设计避免了特征张量在进程间反复拷贝是规模化训练时的关键性能点。内置数据集gb.BuiltinDataset除cora外还内置ogbn-mag、ogbl-citation2、ogbn-arxiv、ogbn-papers100M、ogbn-products、ogb-lsc-mag240m、igb-hom系列、igb-het系列等完整清单见 python/dgl/graphbolt/impl/ondisk_dataset.py覆盖同质/异质、节点分类/链接预测等常见场景且大部分大规模数据集在预处理时加入了反向边并去重。GPU 训练特征与图的固定内存Pinned Memory两个示例都包含一段 GPU 专属优化代码if device torch.device(cuda:0): dataset.graph.pin_memory_() dataset.feature.pin_memory_()将图结构与特征固定到内存pinned memory后GPU 可以通过 UVAUnified Virtual Addressing直接访问它们从而让采样与特征抓取在 GPU 端完成、无需逐批拷贝。配合NeighborSampler的overlap_fetch用独立 CUDA 流重叠图抓取与其他运算与num_gpu_cached_edgesGPU 缓存高频访问的顶点邻域降低 PCIe 带宽压力参数可进一步压榨 GPU 数据管线性能。如果使用 CPU 训练这两行会自然跳过。小结一条管道吃透两类任务对比两个示例可以发现 GraphBolt 的统一抽象节点分类与链接预测的差异最终只体现在数据管道的两个节点上——链接预测多接一个sample_uniform_negative负采样和一个exclude_seed_edges防泄露模型侧用compacted_seeds取节点对。其余环节ItemSampler 切批、copy_to 搬运、sample_neighbor 采样、fetch_feature 抓特征、DataLoader 迭代完全一致。这正是 GraphBolt 的设计哲学把 GNN 训练中高频复用的数据侧逻辑沉淀为可组合的标准件让开发者从手写采样循环 特征索引 batch 拼接中解放出来把精力集中在模型与实验本身。若需深入阅读可继续探索邻居采样实现python/dgl/graphbolt/impl/neighbor_sampler.py负采样实现python/dgl/graphbolt/impl/uniform_negative_sampler.py特征抓取python/dgl/graphbolt/feature_fetcher.py数据加载器python/dgl/graphbolt/dataloader.py赞分享人工智能机器学习深度学习图计算【免费下载链接】dglPython package built to ease deep learning on graph, on top of existing DL frameworks.项目地址https://gitcode.com/gh_mirrors/dg/dgl点击查看免费下载相关推荐DGL GraphBolt 数据加载框架完整指南从 Dataset 到 DataLoader 的模块化 GNN 数据管线DGL GraphBolt 数据加载框架完整指南从 Dataset 到 DataLoader 的模块化 GNN 数据管线 导读 dgl.graphbolt 是人工智能机器学习深度学习图计算使用 GraphBolt 开启 GNN 随机训练DGL 可扩展数据加载流水线框架深度解析使用 GraphBolt 开启 GNN 随机训练DGL 可扩展数据加载流水线框架深度解析 GraphBolt 是 DGL 生态中面向图神经网络GNN随机训人工智能机器学习深度学习图计算数据管道构建PyTorch DataLoader高效使用数据管道构建PyTorch DataLoader高效使用 还在为深度学习项目中的数据加载效率低下而烦恼还在手动处理数据批处理和内存管理本文将深入解析PyT示例工程教程创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表