劳工采样器

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

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

Parameters:

  • fanouts (list[int] or list[dict[etype, int]]) –

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

    如果只提供一个整数，DGL会假设每种边类型都有相同的fanout。

    如果某一层的某一类型边提供了-1，则该类型边的所有入边都将被包含。

  • edge_dir (str, 默认 'in') – 可以是 'in'，其中邻居将根据传入边进行采样，或者 'out'，与 dgl.sampling.sample_neighbors() 相同。

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

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

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

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

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

  • prefetch_labels (list[str] or dict[ntype, list[str]], optional) – 为最后一个MFG预取的目标节点数据，对应于小批量的节点标签。

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

  • output_device (device, optional) – The device of the output subgraphs or MFGs. Default is the same as the minibatch of seed nodes.

示例

节点分类

为了在一组节点上训练一个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)

如果您希望进行非均匀劳动力采样：

>>> # 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() 以获取更多详细信息。

注释

关于MFGs的概念，请参考 用户指南第6节 和 小批量训练教程。