DataLoader

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 封装成一个可迭代对象，用于遍历样本的 mini-batch。

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

参数:

  • graph (DGLGraph) – 图。

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

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

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

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

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

    每次迭代生成的 MFG 的设备，应为 PyTorch 设备对象（例如 torch.device）。

    默认值为 None。如果 use_uva 为 True，MFG 和图将在 torch.cuda.current_device() 中生成，否则将在 g 的相同设备中生成。

  • use_ddp (布尔值, 可选) –

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

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

  • ddp_seed (int, 可选) –

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

    仅当 use_ddp 为 True 时有效。

  • use_uva (bool, 可选) –

    是否使用 Unified Virtual Addressing (UVA) 直接从 CPU 采样图并将特征切片到 GPU。将其设置为 True 将把图和特征张量固定到锁定内存中。

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

    默认值：False。

  • use_prefetch_thread (bool, 可选) –

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

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

  • use_alternate_streams (bool, 可选) –

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

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

  • pin_prefetcher (bool, 可选) –

    （高级选项）是否将特征张量固定到锁定内存中。

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

  • gpu_cache (dict[dict], 可选) –

    使用 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): 是否在每个 epoch 中随机打乱索引。

示例

在同构图上，对一组节点 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)

注意

有关用法，请参阅 Minibatch 训练教程 和 用户指南第 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 上进行。