LaborSampler

class dgl.dataloading.LaborSampler(fanouts, edge_dir='in', prob=None, importance_sampling=0, layer_dependency=False, batch_dependency=1, prefetch_node_feats=None, prefetch_labels=None, prefetch_edge_feats=None, output_device=None)[source]

基类：BlockSampler

一种采样器，用于构建多层 GNN 的节点表示计算依赖关系，通过来自 NeurIPS 2023 论文 Layer-Neighbor Sampling – Defusing Neighborhood Explosion in GNNs 的 Labor 采样方法实现。

此采样器将使每个节点从每种边类型的固定数量的邻居收集消息。默认参数下，邻居是均匀选择的。对于每个将被考虑采样的顶点 t，将存在一个单一的随机变量 r_t。

参数:

  • fanouts (list[int] 或 list[dict[etype, int]])

    一个列表，指定每个 GNN 层每种边类型采样的邻居数量，其中第 i 个元素是第 i 个 GNN 层的 fanout。

    如果只提供一个整数，DGL 假定每种边类型将具有相同的 fanout。

    如果某个层上某个边类型的值为 -1，则包含该边类型的所有入边。

  • edge_dir (str, 默认值 'in') – 可以是 'in'，表示根据入边采样邻居；或者 'out'，表示根据出边采样，与 dgl.sampling.sample_neighbors() 相同。

  • prob (str, 可选) – 如果给定，每个邻居被采样的概率与 g.edata 中指定名称的边特征值成比例。该特征必须是每条边上的标量值。在这种情况下，返回的 block 的 edata 将包含 'edge_weights'，需要在消息传递操作中使用。

  • importance_sampling (int, 默认值 0) – 是否使用重要性采样或均匀采样。使用负值会优化重要性采样概率直到收敛，而使用正值会运行多次优化步骤。如果值为 i，则使用 LABOR-i 变体。当使用非零参数时，返回的 block 的 edata 将包含 'edge_weights'，需要在消息传递操作中使用。

  • layer_dependency (bool, 默认值 False) – 指定不同的层是否应使用相同的随机变量。这将减少采样的顶点数量，但可能会轻微降低质量。

  • batch_dependency (int, 默认值 1) – 指定不同的 minibatch 是否应使用相似的随机变量。这将提高采样顶点的时间访问局部性，但可能会轻微降低质量。

  • prefetch_node_feats (list[str] 或 dict[ntype, list[str]], 可选) – 为第一个 MFG 预取的源节点数据，对应于第一个 GNN 层所需的输入节点特征。

  • prefetch_labels (list[str] 或 dict[ntype, list[str]], 可选) – 为最后一个 MFG 预取的目标节点数据，对应于 minibatch 的节点标签。

  • prefetch_edge_feats (list[str] 或 dict[etype, list[str]], 可选) – 为所有 MFG 预取的边数据名称，对应于所有 GNN 层所需的边特征。

  • output_device (device, 可选) – 输出子图或 MFG 的设备。默认为与种子节点的 minibatch 相同的设备。

示例

节点分类

在同构图上训练一个用于节点分类的 3 层 GNN，节点集为 train_nid，其中每个节点在第一、第二和第三层分别从 5、10、15 个邻居接收消息（假设后端是 PyTorch）

>>> sampler = dgl.dataloading.LaborSampler([5, 10, 15])
>>> dataloader = dgl.dataloading.DataLoader(
...     g, train_nid, sampler,
...     batch_size=1024, shuffle=True, drop_last=False, num_workers=4)
>>> for input_nodes, output_nodes, blocks in dataloader:
...     train_on(blocks)

如果在异构图上进行训练，并且希望每种边类型有不同的邻居数量，则应改为提供一个字典列表。每个字典指定每种边类型要选择的邻居数量。

>>> sampler = dgl.dataloading.LaborSampler([
...     {('user', 'follows', 'user'): 5,
...      ('user', 'plays', 'game'): 4,
...      ('game', 'played-by', 'user'): 3}] * 3)

如果希望使用非均匀 Labor 采样

>>> # any non-negative 1D vector works
>>> g.edata['p'] = torch.rand(g.num_edges())
>>> sampler = dgl.dataloading.LaborSampler([5, 10, 15], prob='p')

边分类和链接预测

此类别也可以与 as_edge_prediction_sampler() 一起用于边分类和链接预测。

>>> sampler = dgl.dataloading.LaborSampler([5, 10, 15])
>>> sampler = dgl.dataloading.as_edge_prediction_sampler(sampler)
>>> dataloader = dgl.dataloading.DataLoader(
...     g, train_eid, sampler,
...     batch_size=1024, shuffle=True, drop_last=False, num_workers=4)

有关更多详细信息，请参阅文档 as_edge_prediction_sampler()。

注意事项

有关 MFG 的概念，请参阅用户指南第 6 节和Minibatch 训练教程。