6.8 使用 GPU 进行邻居采样

注意

GraphBolt 尚未支持基于 GPU 的邻居采样。因此本指南使用 DataLoader 进行说明。

DGL 从 0.7 版本开始支持基于 GPU 的邻居采样，这比基于 CPU 的邻居采样具有显著的速度优势。如果你估计你的图可以放入 GPU 内存中，并且你的模型不占用大量 GPU 内存，那么最好将图放入 GPU 内存中并使用基于 GPU 的邻居采样。

例如，OGB Products 数据集有 240 万个节点和 6100 万条边。由于图的内存消耗取决于边的数量，因此该图占用的内存少于 1GB。所以将整个图放入 GPU 中是完全可能的。

在 DGL 数据加载器中使用基于 GPU 的邻居采样

可以通过以下方式在 DGL 数据加载器中使用基于 GPU 的邻居采样：

  • 将图放入 GPU。

  • 将 train_nid 放入 GPU。

  • 将 device 参数设置为 GPU 设备。

  • 将 num_workers 参数设置为 0，因为 CUDA 不允许多个进程访问同一个上下文。

用于 DataLoader 的所有其他参数可以与其他用户指南和教程中的设置相同。

g = g.to('cuda:0')
train_nid = train_nid.to('cuda:0')
dataloader = dgl.dataloading.DataLoader(
    g,                                # The graph must be on GPU.
    train_nid,                        # train_nid must be on GPU.
    sampler,
    device=torch.device('cuda:0'),    # The device argument must be GPU.
    num_workers=0,                    # Number of workers must be 0.
    batch_size=1000,
    drop_last=False,
    shuffle=True)

注意

基于 GPU 的邻居采样也适用于自定义邻居采样器，前提是 (1) 你的采样器继承自 BlockSampler，并且 (2) 你的采样器完全在 GPU 上运行。

在 DGL 数据加载器中使用基于 CUDA UVA 的邻居采样

注意

DGL 0.8 版本中引入的新特性。

对于图太大无法放入 GPU 内存的情况，我们引入了基于 CUDA UVA (Unified Virtual Addressing) 的采样，其中 GPU 通过零拷贝访问在固定于 CPU 内存中的图上执行采样。您可以通过以下方式在 DGL 数据加载器中启用基于 UVA 的邻居采样：

  • 将 train_nid 放入 GPU。

  • 将 device 参数设置为 GPU 设备。

  • 将 num_workers 参数设置为 0，因为 CUDA 不允许多个进程访问同一个上下文。

  • 设置 use_uva=True。

用于 DataLoader 的所有其他参数可以与其他用户指南和教程中的设置相同。

train_nid = train_nid.to('cuda:0')
dataloader = dgl.dataloading.DataLoader(
    g,
    train_nid,                        # train_nid must be on GPU.
    sampler,
    device=torch.device('cuda:0'),    # The device argument must be GPU.
    num_workers=0,                    # Number of workers must be 0.
    batch_size=1000,
    drop_last=False,
    shuffle=True,
    use_uva=True)                     # Set use_uva=True

基于 UVA 的采样是大图上进行 mini-batch 训练的推荐解决方案，特别适用于多 GPU 训练。

注意

要在多 GPU 训练中使用基于 UVA 的采样，应在启动训练进程之前先将图的所有必要稀疏格式进行实例化。有关更多详细信息，请参阅我们的 GraphSAGE 示例。

对 PinSAGESampler/RandomWalkNeighborSampler 的 UVA 和 GPU 支持

PinSAGESampler 和 RandomWalkNeighborSampler 支持 UVA 和 GPU 采样。您可以通过以下方式启用它们：

  • 固定图（用于 UVA 采样）或将图放入 GPU（用于 GPU 采样）。

  • 将 train_nid 放入 GPU。

g = dgl.heterograph({
    ('item', 'bought-by', 'user'): ([0, 0, 1, 1, 2, 2, 3, 3], [0, 1, 0, 1, 2, 3, 2, 3]),
    ('user', 'bought', 'item'): ([0, 1, 0, 1, 2, 3, 2, 3], [0, 0, 1, 1, 2, 2, 3, 3])})

# UVA setup
# g.create_formats_()
# g.pin_memory_()

# GPU setup
device = torch.device('cuda:0')
g = g.to(device)

sampler1 = dgl.sampling.PinSAGESampler(g, 'item', 'user', 4, 0.5, 3, 2)
sampler2 = dgl.sampling.RandomWalkNeighborSampler(g, 4, 0.5, 3, 2, ['bought-by', 'bought'])

train_nid = torch.tensor([0, 2], dtype=g.idtype, device=device)
sampler1(train_nid)
sampler2(train_nid)

使用 DGL 函数进行基于 GPU 的邻居采样

您可以使用以下支持在 GPU 上操作的函数构建自己的 GPU 采样管道：

子图提取操作

用于子图构建的图转换操作