PyG 邻居采样Neighbor Sampling实战指南用 NeighborLoader 将 GNN 扩展到大规模图【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric导读本文围绕 PyTorch GeometricPyG官方教程 docs/source/tutorial/neighbor_loader.rst 展开深入讲解如何借助torch_geometric.loader.NeighborLoader以节点级node-wise邻居采样方式对大图进行 mini-batch 训练从而化解 GNN 规模化中棘手的neighbor explosion邻居爆炸问题。读完本文你将掌握 NeighborLoader 的核心参数语义、采样输出子图的解析方法、与 GraphSAGE 等模型组合的完整训练循环以及面向异构图、时间图、链接预测等场景的高级用法并了解其底层采样实现原理。为什么 GNN 难以做 mini-batch 训练传统深度神经网络之所以能轻松扩展到海量数据是因为其损失函数可以分解为独立样本mini-batch从而用随机梯度近似精确梯度。但 GNN 面临一个根本性困难单个节点的嵌入递归地依赖于它所有邻居的嵌入这种节点间的相互依赖会随着层数增加呈指数级增长也就是著名的neighbor explosion现象。作为简单的变通方案GNN 通常以 full-batch全图方式训练参见 examples/gcn.py即 GNN 在每一层都能访问所有节点的隐藏表示。然而在工业级、社交网络级的大规模图上这受限于内存容量且收敛缓慢并不可行。因此将 GNN 应用于大规模图必须借助可扩展性技术来缓解 mini-batch 训练引发的邻居爆炸问题主要包括节点级node-wise采样如邻居采样本文的核心主题层间layer-wise采样逐层控制参与消息传递的节点集合子图级subgraph-wise采样如 ClusterLoader、GraphSAINTSampler将传播与预测解耦decouple propagations from predictions如 CorrectAndSmooth、SIGN 等。本文聚焦最常用的节点级采样——邻居采样其思想最早由 Hamilton 等人在论文Inductive Representation Learning on Large GraphsarXiv:1706.02216中提出。NeighborLoader 与邻居采样的核心原理PyG 通过torch_geometric.loader.NeighborLoader类实现邻居采样源码。其工作方式如下对每个节点v ∈ V递归地采样至多k个邻居即Ñ(v) ⊂ N(v)且|Ñ(v)| ≤ k从而使得整体的 L 跳邻域大小被限制在O(k^L)。具体过程是从一组种子节点B ⊂ V出发为B中每个节点采样至多k个邻居然后对上一跳采样到的每个节点再次采样邻居如此递归进行。最终得到的图结构是围绕每个种子节点v ∈ B的一个有向 L 跳子图且保证子图中每个节点到至少一个种子节点存在一条长度不超过 L 的路径。因此一个 L 层的消息传递 GNN 将把采样到的全部节点纳入其计算图。需要清醒认识的是邻居采样只能在一定程度上缓解邻居爆炸——整体邻域规模仍然随层数指数增长因此一般采样超过两到三跳就不可行了。采样跳数与 GNN 层数的对齐实践中采样跳数通常与消息传递层数保持同步。原因很直观如果采样跳数多于消息传递层数GNN 永远无法把后几跳采样到的节点特征融入种子节点的最终表示白白浪费计算资源。当然也可以使用更深的 GNN但此时必须将采样得到的子图转换为双向形式以保证消息传递流向正确。PyG 通过 NeighborLoader 的一个额外参数subgraph_type支持这一点而其他 mini-batch 技术如 ClusterLoader、GraphSAINTSampler、ShaDowKHopSampler 则天然为这类场景设计。基本用法从零构建一个采样器NeighborLoader 从 PyG 的 Data 或 HeteroData 对象初始化并定义如何执行采样。其关键参数包括参数含义input_nodes定义开始采样的种子节点集合为None时使用全部节点num_neighbors定义每一跳为每个节点采样的邻居数量某条目设为-1表示包含全部邻居batch_size定义每次考虑的种子节点数量继承自torch.utils.data.DataLoaderreplace是否放回采样有放回/无放回shuffle每个 epoch 是否对种子节点打乱下面是教程中的完整示例import torch from torch_geometric.data import Data from torch_geometric.loader import NeighborLoader x torch.randn(8, 32) # Node features of shape [num_nodes, num_features] y torch.randint(0, 4, (8, )) # Node labels of shape [num_nodes] edge_index torch.tensor([ [2, 3, 3, 4, 5, 6, 7], [0, 0, 1, 1, 2, 3, 4]], ) # 0 1 # / \/ \ # 2 3 4 # | | | # 5 6 7 data Data(xx, yy, edge_indexedge_index) loader NeighborLoader( data, input_nodestorch.tensor([0, 1]), num_neighbors[2, 1], batch_size1, replaceFalse, shuffleFalse, )这里我们为前两个节点采样子图第一跳采样 2 个邻居第二跳采样 1 个邻居batch_size1会把input_nodes切分成大小为 1 的块。预期行为是种子节点0在第一跳采样到节点2和3第二跳中节点2采样到5节点3采样到6。让我们验证一下batch next(iter(loader)) batch.edge_index tensor([[1, 2, 3, 4], [0, 0, 1, 2]]) batch.n_id tensor([0, 2, 3, 5, 6]) batch.batch_size 1理解返回的 mini-batch 结构NeighborLoader 返回一个 Data 对象包含以下关键属性batch.edge_index采样子图的边索引batch.n_id所有采样节点的原始全局节点索引batch.batch_size种子节点数量即 batch 大小。此外节点特征和边特征会被自动过滤只保留被采样节点/边对应的部分。这一点由底层 NodeLoader 的filter_fn完成——它将采样输出SamplerOutput与原始特征缝合成新的 Data 对象。注意batch.edge_index中的节点索引已被重标号范围是0到batch.num_nodes - 1。要还原原始节点索引只需batch.n_id[batch.edge_index] tensor([[2, 3, 5, 6], [0, 0, 2, 3]])此外还有两个值得注意的细节边方向与消息传递一致NeighborLoader 从种子节点出发采样但返回的子图中的边是指向种子节点的从源到目的这恰好与 PyG 默认的源→目的消息传递流程一致无需任何额外处理即可直接送入 GNN 层。节点排序保证输出节点保证有序前batch_size个采样节点恰好就是用于采样的种子节点batch.n_id[:batch.batch_size] tensor([0])更多输出属性源码级补充除教程列出的三个属性外从 node_loader.py 的实现可以看到返回的 mini-batch 还会附带e_id每条采样边对应的全局边索引用于按边索引取回边特征input_idinput_nodes的全局索引num_sampled_nodes每一跳采样到的节点数num_sampled_edges每一跳采样到的边数seed_time在时间采样场景下的种子时间戳。这些元数据在调试采样结果、做逐层推理见下文 Reddit 示例时非常有用。用 NeighborLoader 训练 GNN完整训练循环得到采样器后就可以把它当作标准的数据加载流程来训练大规模图上的 GNN。首先构造一个简单的两层 GraphSAGE 模型from torch_geometric.nn import GraphSAGE device torch.device(cuda if torch.cuda.is_available() else cpu) model GraphSAGE( in_channels32, hidden_channels64, out_channels4, num_layers2 ).to(device) optimizer torch.optim.Adam(model.parameters(), lr0.01)然后将loader与model结合成训练循环import torch.nn.functional as F for batch in loader: optimizer.zero_grad() batch batch.to(device) out model(batch.x, batch.edge_index) # NOTE Only consider predictions and labels of seed nodes: y batch.y[:batch.batch_size] out out[:batch.batch_size] loss F.cross_entropy(out, y) loss.backward() optimizer.step()这个训练循环与任何标准 PyTorch 训练循环别无二致唯一重要的区别是模型默认会输出形状为[batch.num_nodes, *]的矩阵而我们只关心种子节点的预测。因此通过高效切片同时作用于预测out和标签batch.y只取前batch_size个节点参与损失与指标计算确保梯度只由真实种子节点的预测驱动。底层实现NeighborLoader 如何工作从源码结构看NeighborLoader 的职责被清晰分层neighbor_loader.pyNeighborLoader 初始化时若未显式传入neighbor_sampler会内部构造一个 NeighborSampler把num_neighbors、replace、subgraph_type、disjoint、time_attr、weight_attr等参数传递给它NeighborSampler 在初始化时把图转换为CSC列压缩格式同构图用to_csc异构图用to_hetero_csc为逐跳采样做好准备每次迭代NodeLoader 的collate_fn取出当前批次的种子节点调用node_sampler.sample_from_nodes真正的采样发生在NeighborSampler._sample中优先调用pyg-lib提供的torch.ops.pyg.neighbor_sample异构图对应hetero_neighbor_sample若未安装则回退到torch-sparse的torch.ops.torch_sparse.neighbor_sample。两个后端都未安装时会抛出ImportError——这是使用 NeighborLoader 的硬性依赖前提neighbor_sampler.py采样结果SamplerOutput定义于 sampler/base.py包含局部重标号的row/col、全局节点索引node、全局边索引edge等再由 NodeLoader 的filter_fn过滤特征并组装成最终的 Data mini-batch。关于后端还有一个值得注意的提示在 Linux 上使用非induced子图类型时若未安装pyg-libNeighborSampler 会给出弃用警告建议安装pyg-lib以获得加速的邻居采样neighbor_sampler.py。参数详解来自源码的完整说明从 NeighborLoader 构造函数 及 NeighborSampler 的签名我们可以整理出更完整的参数语义input_time为input_nodes覆写时间戳设置它时必须同时设置time_attr否则抛出ValueErrortime_attr指定图上节点级或边级的时间戳属性。一旦设置将启用时间感知采样保证邻居的时间戳不晚于中心节点支持节点级与边级两种时间源码会做严格校验temporal_strategy时间采样策略uniform在满足时间约束的邻居中均匀采样默认或last取满足约束的最后num_neighbors个邻居weight_attr指定边权重属性以启用加权采样权重越高的邻居越容易被采样。权重不需要归一化但必须非负、有限且局部邻域内和不为零边级时间采样与加权采样分别要求较新的pyg-lib版本neighbor_sampler.pyis_sorted若edge_index已按列排序可置为True以跳过内部重排序、提升运行时与内存效率filter_per_worker控制特征过滤发生在 worker 子进程还是主进程None时依据数据是否部分驻留 GPU 自动推断transform/transform_sampler_output分别对采样后的 mini-batch 或原始采样输出做后处理其余**kwargs全部透传给torch.utils.data.DataLoader如num_workers、drop_last等。这些高级参数均有对应的测试用例覆盖例如 test/loader/test_neighbor_loader.py 中对weight_attr第 836、869 行附近与time_attr第 899、930 行附近的验证。大规模实战Reddit 示例剖析教程指出一个在真实大规模数据上可运行的完整示例位于 examples/reddit.py。该示例的工程实践非常值得借鉴kwargs {batch_size: 1024, num_workers: 6, persistent_workers: True} train_loader NeighborLoader(data, input_nodesdata.train_mask, num_neighbors[25, 10], shuffleTrue, **kwargs) subgraph_loader NeighborLoader(copy.copy(data), input_nodesNone, num_neighbors[-1], shuffleFalse, **kwargs)它展示了三个关键技巧训练用采样器num_neighbors[25, 10]表示两跳分别采样 25 和 10 个邻居input_nodesdata.train_mask只从训练节点出发采样评估用全邻居采样器num_neighbors[-1]表示包含全部邻居-1语义并用copy.copy(data)避免污染原图评估时删除x/y特征并手动注入全局节点索引n_id以节省内存逐层推理layer-wise inference模型类中的inference方法不一次性前向整个 batch而是逐层对全图所有节点计算表示——每一层都用x_all[batch.n_id]取回上一层的嵌入通过subgraph_loader分批完成最后仅保留种子节点的表示x[:batch.batch_size]。这避免了深层堆叠导致的大内存占用。分层扩展消除冗余计算NeighborLoader 的一个固有缺点是它会为所有采样节点在所有网络深度上计算表示。然而后几跳采样到的节点在后面的 GNN 层中已不再贡献于种子节点的表示属于无用计算会让 NeighborLoader 略微变慢。这是为获得干净、模块化、便于实验的 GNN 设计而做出的取舍——模型定义不与数据加载方式耦合。如果希望消除这一开销、进一步加速 mini-batch GNN 的训练与推理可以参考 分层邻居采样Hierarchical Neighborhood Sampling教程其中介绍了如何利用逐层采样的计算图来剪枝冗余前向传播。高级选项异构、分离子图与更深的 GNN同构与异构图的开箱即用支持NeighborLoader 开箱即用地同时支持同构图与异构图采样只要传入 HeteroData 对象即可。在异构图上采样支持对采样参数进行细粒度控制——例如可以为每种边类型单独指定采样邻居数源码文档字符串loader NeighborLoader( hetero_data, # Sample 30 neighbors for each node and edge type for 2 iterations num_neighbors{key: [30] * 2 for key in hetero_data.edge_types}, # Use a batch size of 128 for sampling training nodes of type paper batch_size128, input_nodes(paper, hetero_data[paper].train_mask), )这里input_nodes需以(节点类型, 掩码)元组形式指定num_neighbors则可以是按边类型分组的字典。同构与异构的完整示例可分别参考 examples/reddit.py 与 examples/hetero/to_hetero_mag.py。disjoint是否为每个种子节点构建独立子图默认情况下NeighborLoader 会把不同种子节点采样到的节点融合进同一个子图这样共享邻居不会被重复采样从而节省内存。可以通过传入disjointTrue禁用该行为——此时每个种子节点拥有自己独立的子图输出中会附带一个batch向量用于标识每个节点属于哪个子图。源码中还有一个值得注意的细节时间采样会自动强制disjointTrueneighbor_sampler.py因为不同时间步的种子节点必须保持各自的邻域独立性。subgraph_type为更深 GNN 定制子图结构默认返回的子图是**有向directional**的这限制了它只能用于层数与采样跳数相等的 GNN。若想使用更深的 GNN可通过subgraph_type参数调整directional默认仅保留计算种子节点表示所必需的采样有向边bidirectional将采样到的边转换为双向边induced返回所有采样节点构成的诱导子图即包含采样节点之间的全部边。其中SubgraphType枚举定义于 sampler/base.pybidirectional的转换实现在SamplerOutput.to_bidirectional中。注意历史参数directed已弃用——从 neighbor_sampler.py 可以看到directedFalse会被自动映射为subgraph_typeinduced并发出弃用警告新代码请直接使用subgraph_type。链接预测场景LinkNeighborLoaderNeighborLoader 面向从单个种子节点出发采样设计因此不直接适用于链接预测。对于链接预测场景PyG 提供了 LinkNeighborLoader它接收一组输入边并从源节点和目的节点两侧同时进行邻居采样来构造子图。总结邻居采样是 GNN 规模化最常用、最基础的技巧之一。通过本文你可以看到问题本质邻居爆炸使得朴素 mini-batch 训练在大图上不可行而邻居采样通过限制每跳邻居数num_neighbors把子图规模控制在O(k^L)核心工具NeighborLoader 一行即可完成种子节点选取 → 逐跳采样 → 特征过滤 → mini-batch 组装全流程返回的 Data 对象通过n_id/e_id/batch_size等属性与全局图保持可追溯关系训练范式只需在训练循环中把预测与标签切片到前batch_size个种子节点即可复用标准 PyTorch 训练流程进阶能力同构/异构、时间感知采样、加权采样、disjoint、subgraph_type、逐层推理等特性让 NeighborLoader 足以覆盖绝大多数大规模图学习场景原理纵深从源码看采样真正执行于 NeighborSampler依赖pyg-lib或torch-sparse后端与 NodeLoader 的过滤逻辑配合完成整个数据管线。当你下一次面对百万甚至十亿级节点的图数据时NeighborLoader 会是你进入 mini-batch 训练世界的第一个、也是最重要的入口。【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 
长河君 · 2026/9/12 21:04:02 · 阅读 1.2k PyG 邻居采样（Neighbor Sampling）实战指南：用 NeighborLoader 将 GNN 扩展到大规模图
PyG 邻居采样Neighbor Sampling实战指南用 NeighborLoader 将 GNN 扩展到大规模图【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric导读本文围绕 PyTorch GeometricPyG官方教程 docs/source/tutorial/neighbor_loader.rst 展开深入讲解如何借助torch_geometric.loader.NeighborLoader以节点级node-wise邻居采样方式对大图进行 mini-batch 训练从而化解 GNN 规模化中棘手的neighbor explosion邻居爆炸问题。读完本文你将掌握 NeighborLoader 的核心参数语义、采样输出子图的解析方法、与 GraphSAGE 等模型组合的完整训练循环以及面向异构图、时间图、链接预测等场景的高级用法并了解其底层采样实现原理。为什么 GNN 难以做 mini-batch 训练传统深度神经网络之所以能轻松扩展到海量数据是因为其损失函数可以分解为独立样本mini-batch从而用随机梯度近似精确梯度。但 GNN 面临一个根本性困难单个节点的嵌入递归地依赖于它所有邻居的嵌入这种节点间的相互依赖会随着层数增加呈指数级增长也就是著名的neighbor explosion现象。作为简单的变通方案GNN 通常以 full-batch全图方式训练参见 examples/gcn.py即 GNN 在每一层都能访问所有节点的隐藏表示。然而在工业级、社交网络级的大规模图上这受限于内存容量且收敛缓慢并不可行。因此将 GNN 应用于大规模图必须借助可扩展性技术来缓解 mini-batch 训练引发的邻居爆炸问题主要包括节点级node-wise采样如邻居采样本文的核心主题层间layer-wise采样逐层控制参与消息传递的节点集合子图级subgraph-wise采样如 ClusterLoader、GraphSAINTSampler将传播与预测解耦decouple propagations from predictions如 CorrectAndSmooth、SIGN 等。本文聚焦最常用的节点级采样——邻居采样其思想最早由 Hamilton 等人在论文Inductive Representation Learning on Large GraphsarXiv:1706.02216中提出。NeighborLoader 与邻居采样的核心原理PyG 通过torch_geometric.loader.NeighborLoader类实现邻居采样源码。其工作方式如下对每个节点v ∈ V递归地采样至多k个邻居即Ñ(v) ⊂ N(v)且|Ñ(v)| ≤ k从而使得整体的 L 跳邻域大小被限制在O(k^L)。具体过程是从一组种子节点B ⊂ V出发为B中每个节点采样至多k个邻居然后对上一跳采样到的每个节点再次采样邻居如此递归进行。最终得到的图结构是围绕每个种子节点v ∈ B的一个有向 L 跳子图且保证子图中每个节点到至少一个种子节点存在一条长度不超过 L 的路径。因此一个 L 层的消息传递 GNN 将把采样到的全部节点纳入其计算图。需要清醒认识的是邻居采样只能在一定程度上缓解邻居爆炸——整体邻域规模仍然随层数指数增长因此一般采样超过两到三跳就不可行了。采样跳数与 GNN 层数的对齐实践中采样跳数通常与消息传递层数保持同步。原因很直观如果采样跳数多于消息传递层数GNN 永远无法把后几跳采样到的节点特征融入种子节点的最终表示白白浪费计算资源。当然也可以使用更深的 GNN但此时必须将采样得到的子图转换为双向形式以保证消息传递流向正确。PyG 通过 NeighborLoader 的一个额外参数subgraph_type支持这一点而其他 mini-batch 技术如 ClusterLoader、GraphSAINTSampler、ShaDowKHopSampler 则天然为这类场景设计。基本用法从零构建一个采样器NeighborLoader 从 PyG 的 Data 或 HeteroData 对象初始化并定义如何执行采样。其关键参数包括参数含义input_nodes定义开始采样的种子节点集合为None时使用全部节点num_neighbors定义每一跳为每个节点采样的邻居数量某条目设为-1表示包含全部邻居batch_size定义每次考虑的种子节点数量继承自torch.utils.data.DataLoaderreplace是否放回采样有放回/无放回shuffle每个 epoch 是否对种子节点打乱下面是教程中的完整示例import torch from torch_geometric.data import Data from torch_geometric.loader import NeighborLoader x torch.randn(8, 32) # Node features of shape [num_nodes, num_features] y torch.randint(0, 4, (8, )) # Node labels of shape [num_nodes] edge_index torch.tensor([ [2, 3, 3, 4, 5, 6, 7], [0, 0, 1, 1, 2, 3, 4]], ) # 0 1 # / \/ \ # 2 3 4 # | | | # 5 6 7 data Data(xx, yy, edge_indexedge_index) loader NeighborLoader( data, input_nodestorch.tensor([0, 1]), num_neighbors[2, 1], batch_size1, replaceFalse, shuffleFalse, )这里我们为前两个节点采样子图第一跳采样 2 个邻居第二跳采样 1 个邻居batch_size1会把input_nodes切分成大小为 1 的块。预期行为是种子节点0在第一跳采样到节点2和3第二跳中节点2采样到5节点3采样到6。让我们验证一下batch next(iter(loader)) batch.edge_index tensor([[1, 2, 3, 4], [0, 0, 1, 2]]) batch.n_id tensor([0, 2, 3, 5, 6]) batch.batch_size 1理解返回的 mini-batch 结构NeighborLoader 返回一个 Data 对象包含以下关键属性batch.edge_index采样子图的边索引batch.n_id所有采样节点的原始全局节点索引batch.batch_size种子节点数量即 batch 大小。此外节点特征和边特征会被自动过滤只保留被采样节点/边对应的部分。这一点由底层 NodeLoader 的filter_fn完成——它将采样输出SamplerOutput与原始特征缝合成新的 Data 对象。注意batch.edge_index中的节点索引已被重标号范围是0到batch.num_nodes - 1。要还原原始节点索引只需batch.n_id[batch.edge_index] tensor([[2, 3, 5, 6], [0, 0, 2, 3]])此外还有两个值得注意的细节边方向与消息传递一致NeighborLoader 从种子节点出发采样但返回的子图中的边是指向种子节点的从源到目的这恰好与 PyG 默认的源→目的消息传递流程一致无需任何额外处理即可直接送入 GNN 层。节点排序保证输出节点保证有序前batch_size个采样节点恰好就是用于采样的种子节点batch.n_id[:batch.batch_size] tensor([0])更多输出属性源码级补充除教程列出的三个属性外从 node_loader.py 的实现可以看到返回的 mini-batch 还会附带e_id每条采样边对应的全局边索引用于按边索引取回边特征input_idinput_nodes的全局索引num_sampled_nodes每一跳采样到的节点数num_sampled_edges每一跳采样到的边数seed_time在时间采样场景下的种子时间戳。这些元数据在调试采样结果、做逐层推理见下文 Reddit 示例时非常有用。用 NeighborLoader 训练 GNN完整训练循环得到采样器后就可以把它当作标准的数据加载流程来训练大规模图上的 GNN。首先构造一个简单的两层 GraphSAGE 模型from torch_geometric.nn import GraphSAGE device torch.device(cuda if torch.cuda.is_available() else cpu) model GraphSAGE( in_channels32, hidden_channels64, out_channels4, num_layers2 ).to(device) optimizer torch.optim.Adam(model.parameters(), lr0.01)然后将loader与model结合成训练循环import torch.nn.functional as F for batch in loader: optimizer.zero_grad() batch batch.to(device) out model(batch.x, batch.edge_index) # NOTE Only consider predictions and labels of seed nodes: y batch.y[:batch.batch_size] out out[:batch.batch_size] loss F.cross_entropy(out, y) loss.backward() optimizer.step()这个训练循环与任何标准 PyTorch 训练循环别无二致唯一重要的区别是模型默认会输出形状为[batch.num_nodes, *]的矩阵而我们只关心种子节点的预测。因此通过高效切片同时作用于预测out和标签batch.y只取前batch_size个节点参与损失与指标计算确保梯度只由真实种子节点的预测驱动。底层实现NeighborLoader 如何工作从源码结构看NeighborLoader 的职责被清晰分层neighbor_loader.pyNeighborLoader 初始化时若未显式传入neighbor_sampler会内部构造一个 NeighborSampler把num_neighbors、replace、subgraph_type、disjoint、time_attr、weight_attr等参数传递给它NeighborSampler 在初始化时把图转换为CSC列压缩格式同构图用to_csc异构图用to_hetero_csc为逐跳采样做好准备每次迭代NodeLoader 的collate_fn取出当前批次的种子节点调用node_sampler.sample_from_nodes真正的采样发生在NeighborSampler._sample中优先调用pyg-lib提供的torch.ops.pyg.neighbor_sample异构图对应hetero_neighbor_sample若未安装则回退到torch-sparse的torch.ops.torch_sparse.neighbor_sample。两个后端都未安装时会抛出ImportError——这是使用 NeighborLoader 的硬性依赖前提neighbor_sampler.py采样结果SamplerOutput定义于 sampler/base.py包含局部重标号的row/col、全局节点索引node、全局边索引edge等再由 NodeLoader 的filter_fn过滤特征并组装成最终的 Data mini-batch。关于后端还有一个值得注意的提示在 Linux 上使用非induced子图类型时若未安装pyg-libNeighborSampler 会给出弃用警告建议安装pyg-lib以获得加速的邻居采样neighbor_sampler.py。参数详解来自源码的完整说明从 NeighborLoader 构造函数 及 NeighborSampler 的签名我们可以整理出更完整的参数语义input_time为input_nodes覆写时间戳设置它时必须同时设置time_attr否则抛出ValueErrortime_attr指定图上节点级或边级的时间戳属性。一旦设置将启用时间感知采样保证邻居的时间戳不晚于中心节点支持节点级与边级两种时间源码会做严格校验temporal_strategy时间采样策略uniform在满足时间约束的邻居中均匀采样默认或last取满足约束的最后num_neighbors个邻居weight_attr指定边权重属性以启用加权采样权重越高的邻居越容易被采样。权重不需要归一化但必须非负、有限且局部邻域内和不为零边级时间采样与加权采样分别要求较新的pyg-lib版本neighbor_sampler.pyis_sorted若edge_index已按列排序可置为True以跳过内部重排序、提升运行时与内存效率filter_per_worker控制特征过滤发生在 worker 子进程还是主进程None时依据数据是否部分驻留 GPU 自动推断transform/transform_sampler_output分别对采样后的 mini-batch 或原始采样输出做后处理其余**kwargs全部透传给torch.utils.data.DataLoader如num_workers、drop_last等。这些高级参数均有对应的测试用例覆盖例如 test/loader/test_neighbor_loader.py 中对weight_attr第 836、869 行附近与time_attr第 899、930 行附近的验证。大规模实战Reddit 示例剖析教程指出一个在真实大规模数据上可运行的完整示例位于 examples/reddit.py。该示例的工程实践非常值得借鉴kwargs {batch_size: 1024, num_workers: 6, persistent_workers: True} train_loader NeighborLoader(data, input_nodesdata.train_mask, num_neighbors[25, 10], shuffleTrue, **kwargs) subgraph_loader NeighborLoader(copy.copy(data), input_nodesNone, num_neighbors[-1], shuffleFalse, **kwargs)它展示了三个关键技巧训练用采样器num_neighbors[25, 10]表示两跳分别采样 25 和 10 个邻居input_nodesdata.train_mask只从训练节点出发采样评估用全邻居采样器num_neighbors[-1]表示包含全部邻居-1语义并用copy.copy(data)避免污染原图评估时删除x/y特征并手动注入全局节点索引n_id以节省内存逐层推理layer-wise inference模型类中的inference方法不一次性前向整个 batch而是逐层对全图所有节点计算表示——每一层都用x_all[batch.n_id]取回上一层的嵌入通过subgraph_loader分批完成最后仅保留种子节点的表示x[:batch.batch_size]。这避免了深层堆叠导致的大内存占用。分层扩展消除冗余计算NeighborLoader 的一个固有缺点是它会为所有采样节点在所有网络深度上计算表示。然而后几跳采样到的节点在后面的 GNN 层中已不再贡献于种子节点的表示属于无用计算会让 NeighborLoader 略微变慢。这是为获得干净、模块化、便于实验的 GNN 设计而做出的取舍——模型定义不与数据加载方式耦合。如果希望消除这一开销、进一步加速 mini-batch GNN 的训练与推理可以参考 分层邻居采样Hierarchical Neighborhood Sampling教程其中介绍了如何利用逐层采样的计算图来剪枝冗余前向传播。高级选项异构、分离子图与更深的 GNN同构与异构图的开箱即用支持NeighborLoader 开箱即用地同时支持同构图与异构图采样只要传入 HeteroData 对象即可。在异构图上采样支持对采样参数进行细粒度控制——例如可以为每种边类型单独指定采样邻居数源码文档字符串loader NeighborLoader( hetero_data, # Sample 30 neighbors for each node and edge type for 2 iterations num_neighbors{key: [30] * 2 for key in hetero_data.edge_types}, # Use a batch size of 128 for sampling training nodes of type paper batch_size128, input_nodes(paper, hetero_data[paper].train_mask), )这里input_nodes需以(节点类型, 掩码)元组形式指定num_neighbors则可以是按边类型分组的字典。同构与异构的完整示例可分别参考 examples/reddit.py 与 examples/hetero/to_hetero_mag.py。disjoint是否为每个种子节点构建独立子图默认情况下NeighborLoader 会把不同种子节点采样到的节点融合进同一个子图这样共享邻居不会被重复采样从而节省内存。可以通过传入disjointTrue禁用该行为——此时每个种子节点拥有自己独立的子图输出中会附带一个batch向量用于标识每个节点属于哪个子图。源码中还有一个值得注意的细节时间采样会自动强制disjointTrueneighbor_sampler.py因为不同时间步的种子节点必须保持各自的邻域独立性。subgraph_type为更深 GNN 定制子图结构默认返回的子图是**有向directional**的这限制了它只能用于层数与采样跳数相等的 GNN。若想使用更深的 GNN可通过subgraph_type参数调整directional默认仅保留计算种子节点表示所必需的采样有向边bidirectional将采样到的边转换为双向边induced返回所有采样节点构成的诱导子图即包含采样节点之间的全部边。其中SubgraphType枚举定义于 sampler/base.pybidirectional的转换实现在SamplerOutput.to_bidirectional中。注意历史参数directed已弃用——从 neighbor_sampler.py 可以看到directedFalse会被自动映射为subgraph_typeinduced并发出弃用警告新代码请直接使用subgraph_type。链接预测场景LinkNeighborLoaderNeighborLoader 面向从单个种子节点出发采样设计因此不直接适用于链接预测。对于链接预测场景PyG 提供了 LinkNeighborLoader它接收一组输入边并从源节点和目的节点两侧同时进行邻居采样来构造子图。总结邻居采样是 GNN 规模化最常用、最基础的技巧之一。通过本文你可以看到问题本质邻居爆炸使得朴素 mini-batch 训练在大图上不可行而邻居采样通过限制每跳邻居数num_neighbors把子图规模控制在O(k^L)核心工具NeighborLoader 一行即可完成种子节点选取 → 逐跳采样 → 特征过滤 → mini-batch 组装全流程返回的 Data 对象通过n_id/e_id/batch_size等属性与全局图保持可追溯关系训练范式只需在训练循环中把预测与标签切片到前batch_size个种子节点即可复用标准 PyTorch 训练流程进阶能力同构/异构、时间感知采样、加权采样、disjoint、subgraph_type、逐层推理等特性让 NeighborLoader 足以覆盖绝大多数大规模图学习场景原理纵深从源码看采样真正执行于 NeighborSampler依赖pyg-lib或torch-sparse后端与 NodeLoader 的过滤逻辑配合完成整个数据管线。当你下一次面对百万甚至十亿级节点的图数据时NeighborLoader 会是你进入 mini-batch 训练世界的第一个、也是最重要的入口。【免费下载链接】pytorch_geometricGraph Neural Network Library for PyTorch项目地址: https://gitcode.com/GitHub_Trending/py/pytorch_geometric创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考