6.8 使用GPU进行邻域采样

注意

GraphBolt 目前还不支持基于 GPU 的邻居采样。因此，本指南使用 DataLoader 进行说明。

自0.7版本以来，DGL已经支持基于GPU的邻居采样，与基于CPU的邻居采样相比，具有显著的速度优势。如果你估计你的图可以放入GPU内存，并且你的模型不占用太多GPU内存，那么最好将图放入GPU内存并使用基于GPU的邻居采样。

例如，OGB Products 有 2.4M 个节点和 61M 条边。由于图的内存消耗取决于边的数量，因此该图占用的内存不到 1GB。因此，完全有可能将整个图放入 GPU 中。

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

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

  • 将图放到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（统一虚拟寻址）的采样，其中GPU通过零拷贝访问在固定在CPU内存中的图上执行采样。 您可以通过以下方式在DGL数据加载器中启用基于UVA的邻域采样：

  • 将 train_nid 放到 GPU 上。

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

  • Set num_workers argument to 0, because CUDA does not allow multiple processes accessing the same context.

  • 设置 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的采样是推荐用于大型图上的小批量训练解决方案，特别是对于多GPU训练。

注意

要在多GPU训练中使用基于UVA的采样，您应首先在生成训练过程之前实现图的所有必要稀疏格式。有关更多详细信息，请参阅我们的GraphSAGE示例。

UVA和GPU对PinSAGESampler/RandomWalkNeighborSampler的支持

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

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

  • Put the train_nid onto 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)

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

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

子图提取操作：

用于子图构建的图变换操作：