数据加载器

class dgl.dataloading.DataLoader(graph, indices, graph_sampler, device=None, use_ddp=False, ddp_seed=0, batch_size=1, drop_last=False, shuffle=False, use_prefetch_thread=None, use_alternate_streams=None, pin_prefetcher=None, use_uva=False, gpu_cache=None, **kwargs)[source]

基础类：DataLoader

采样图数据加载器。将DGLGraph和Sampler包装成一个可迭代的小批量样本。

DGL的DataLoader扩展了PyTorch的DataLoader，通过处理图样本的创建和传输。它支持迭代一组节点、边或任何类型的索引，以获取DGLGraph、消息流图（MFGS）或训练图神经网络所需的其他结构的样本。

Parameters:

  • graph (DGLGraph) – The graph.

  • indices (Tensor 或 dict[ntype, Tensor]) –

    索引集合。它可以是整数索引的张量，也可以是类型和索引的字典。

    索引的实际含义由 sample() 方法的 graph_sampler 定义。

  • graph_sampler (dgl.dataloading.Sampler) – 子图采样器。

  • device (设备上下文, 可选) –

    每次迭代中生成的MFGs的设备，应该是一个PyTorch设备对象（例如，torch.device）。

    默认情况下，此值为None。如果use_uva为True，MFGs和图将在torch.cuda.current_device()中生成，否则将在g的同一设备中生成。

  • use_ddp (boolean, optional) –

    如果为True，告诉DataLoader使用torch.utils.data.distributed.DistributedSampler为每个参与进程适当地分割训练集。

    覆盖torch.utils.data.DataLoader的sampler参数。

  • ddp_seed (int, optional) –

    用于在torch.utils.data.distributed.DistributedSampler中打乱数据集的种子。

    仅在use_ddp为True时有效。

  • use_uva (bool, optional) –

    是否使用统一虚拟寻址（UVA）直接从图中采样并将特征从CPU切片到GPU。将其设置为True将把图和特征张量固定在固定内存中。

    如果为True，则要求indices必须与device参数具有相同的设备。

    默认值：False。

  • use_prefetch_thread (bool, optional) –

    (高级选项) 生成一个新的Python线程来异步执行特征切片。这可能会加快速度，但会消耗GPU内存。

    默认值：如果图在CPU上且device是CUDA，则为True。否则为False。

  • use_alternate_streams (bool, optional) –

    (高级选项) 是否在非默认流上切片并将特征传输到GPU。

    默认值：如果图在CPU上，device 是CUDA，并且 use_uva 为False，则为True。否则为False。

  • pin_prefetcher (bool, optional) –

    (高级选项) 是否将特征张量固定在固定内存中。

    默认值：如果图在CPU上且device是CUDA，则为True。否则为False。

  • gpu_cache (dict[dict], optional) –

    使用HugeCTR gpu_cache缓存哪些节点和边特征。示例： {“node”: {“features”: 500000}, “edge”: {“types”: 4000000}} 表示我们希望在GPU缓存中缓存50万个节点的“features”和400万个边的“types”。

    仅在计算能力为70或以上的NVIDIA GPU上支持。 字典包含特征键及其对应的缓存大小。请参阅 https://github.com/NVIDIA-Merlin/HugeCTR/blob/main/gpu_cache/ReadMe.md 以获取更多参考信息。

  • kwargs (dict) –

    要传递给父类 PyTorch torch.utils.data.DataLoader 的关键字参数。常见的参数有：

    • batch_size (int): 每个批次中的索引数量。

    • drop_last (bool): 是否丢弃最后一个不完整的批次。

    • shuffle (bool): 是否在每个周期随机打乱索引。

示例

为了在一组节点train_nid上训练一个3层GNN进行节点分类，在一个同质图上，每个节点在第一层从15个邻居接收消息，第二层从10个邻居接收消息，第三层从5个邻居接收消息（假设后端是PyTorch）：

>>> sampler = dgl.dataloading.MultiLayerNeighborSampler([15, 10, 5])
>>> 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(input_nodes, output_nodes, blocks)

与分布式数据并行一起使用

如果您正在使用PyTorch的分布式训练（例如，当使用 torch.nn.parallel.DistributedDataParallel时），您可以通过开启 use_ddp选项来训练模型：

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

注释

请参考 小批量训练教程 和 用户指南第6节 以了解使用方法。

选择合适设备的提示

  • 如果输入的图 g 在 GPU 上，输出设备 device 必须是相同的 GPU， 并且 num_workers 必须为零。在这种情况下，采样和子图构建将在 GPU 上进行。这是在使用单 GPU 且整个图适合 GPU 内存时的推荐设置。

  • 如果输入图 g 在 CPU 上，而输出设备 device 是 GPU，那么 根据 use_uva 的值：

    • 如果 use_uva 设置为 True，即使 GPU 本身无法容纳整个图，采样和子图构建也将在 GPU 上进行。这是推荐的设置，除非有不支持 UVA 的操作。在这种情况下，num_workers 必须为 0。

    • 否则，采样和子图构建都将在CPU上进行。