6.1 使用邻域采样训练GNN进行节点分类

(中文版)

为了使您的模型能够随机训练，您需要执行以下操作：

  • 定义一个邻域采样器。

  • 为小批量训练调整您的模型。

  • 修改你的训练循环。

以下小节将逐一解决这些步骤。

定义一个邻域采样器和数据加载器

DGL 提供了几个邻居采样器类，这些类根据我们希望计算的节点生成每层所需的计算依赖。

最简单的邻居采样器是NeighborSampler 或等效的函数式接口sample_neighbor()， 它使节点从其邻居处收集消息。

要使用DGL提供的采样器，还需要将其与 DataLoader结合使用，它会在小批量中迭代一组索引（在这种情况下是节点）。

例如，以下代码创建一个DataLoader，它分批遍历ogbn-arxiv的训练节点ID集，并将生成的MFG列表放到GPU上。

import dgl
import dgl.graphbolt as gb
import dgl.nn as dglnn
import torch
import torch.nn as nn
import torch.nn.functional as F

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
dataset = gb.BuiltinDataset("ogbn-arxiv").load()
g = dataset.graph
feature = dataset.feature
train_set = dataset.tasks[0].train_set
datapipe = gb.ItemSampler(train_set, batch_size=1024, shuffle=True)
datapipe = datapipe.sample_neighbor(g, [10, 10]) # 2 layers.
# Or equivalently:
# datapipe = gb.NeighborSampler(datapipe, g, [10, 10])
datapipe = datapipe.fetch_feature(feature, node_feature_keys=["feat"])
datapipe = datapipe.copy_to(device)
dataloader = gb.DataLoader(datapipe)

遍历DataLoader将产生MiniBatch，其中包含一系列特别创建的图，这些图表示每层上的计算依赖关系。为了使用DGL进行训练，你可以通过调用mini_batch.blocks来访问消息流图（MFGs）。

mini_batch = next(iter(dataloader))
print(mini_batch.blocks)

注意

请参阅随机训练教程了解消息流图的概念。

如果您希望开发自己的邻域采样器，或者您想更详细地了解MFGs的概念，请参考 6.4 实现自定义图采样器。

为小批量训练调整你的模型

如果你的消息传递模块都是由DGL提供的，那么将你的模型适应于小批量训练所需的改动是最小的。以一个多层GCN为例。如果你的模型在全图上是如下实现的：

class TwoLayerGCN(nn.Module):
    def __init__(self, in_features, hidden_features, out_features):
        super().__init__()
        self.conv1 = dglnn.GraphConv(in_features, hidden_features)
        self.conv2 = dglnn.GraphConv(hidden_features, out_features)

    def forward(self, g, x):
        x = F.relu(self.conv1(g, x))
        x = F.relu(self.conv2(g, x))
        return x

然后你只需要将g替换为上面生成的blocks。

class StochasticTwoLayerGCN(nn.Module):
    def __init__(self, in_features, hidden_features, out_features):
        super().__init__()
        self.conv1 = dgl.nn.GraphConv(in_features, hidden_features)
        self.conv2 = dgl.nn.GraphConv(hidden_features, out_features)

    def forward(self, blocks, x):
        x = F.relu(self.conv1(blocks[0], x))
        x = F.relu(self.conv2(blocks[1], x))
        return x

上面的DGL GraphConv模块接受由数据加载器生成的blocks中的一个元素作为参数。

每个NN模块的API参考将告诉你 它是否支持接受MFG作为参数。

如果您希望使用自己的消息传递模块，请参考 6.6 实现用于小批量训练的自定义GNN模块。

训练循环

训练循环简单地包括使用定制的批处理迭代器遍历数据集。在每次产生MiniBatch的迭代中，我们：

  1. 通过data.node_features["feat"]访问与输入节点对应的节点特征。这些特征已经被数据加载器移动到目标设备（CPU或GPU）。

  2. 通过data.labels访问与输出节点对应的节点标签。这些标签已经被数据加载器移动到目标设备（CPU或GPU）。

  3. 将MFGs列表和输入节点特征提供给多层GNN并获取输出。

  4. 计算损失并反向传播。

model = StochasticTwoLayerGCN(in_features, hidden_features, out_features)
model = model.to(device)
opt = torch.optim.Adam(model.parameters())

for data in dataloader:
    input_features = data.node_features["feat"]
    output_labels = data.labels
    output_predictions = model(data.blocks, input_features)
    loss = compute_loss(output_labels, output_predictions)
    opt.zero_grad()
    loss.backward()
    opt.step()

DGL 提供了一个端到端的随机训练示例 GraphSAGE 实现。

For heterogeneous graphs

在异质图上训练用于节点分类的图神经网络是类似的。

例如，我们之前已经看到 如何在全图上训练一个2层RGCN。 在minibatch训练中实现RGCN的代码看起来非常相似（为了简化，去除了自环、非线性和基分解）：

class StochasticTwoLayerRGCN(nn.Module):
    def __init__(self, in_feat, hidden_feat, out_feat, rel_names):
        super().__init__()
        self.conv1 = dglnn.HeteroGraphConv({
                rel : dglnn.GraphConv(in_feat, hidden_feat, norm='right')
                for rel in rel_names
            })
        self.conv2 = dglnn.HeteroGraphConv({
                rel : dglnn.GraphConv(hidden_feat, out_feat, norm='right')
                for rel in rel_names
            })

    def forward(self, blocks, x):
        x = self.conv1(blocks[0], x)
        x = self.conv2(blocks[1], x)
        return x

DGL 提供的采样器也支持异构图。 例如，仍然可以使用提供的 NeighborSampler 类和 DataLoader 类进行 随机训练。唯一的区别是现在 itemset 是一个 ItemSetDict 的实例，它是一个 节点类型到节点 ID 的字典。

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
dataset = gb.BuiltinDataset("ogbn-mag").load()
g = dataset.graph
feature = dataset.feature
train_set = dataset.tasks[0].train_set
datapipe = gb.ItemSampler(train_set, batch_size=1024, shuffle=True)
datapipe = datapipe.sample_neighbor(g, [10, 10]) # 2 layers.
# Or equivalently:
# datapipe = gb.NeighborSampler(datapipe, g, [10, 10])
# For heterogeneous graphs, we need to specify the node feature keys
# for each node type.
datapipe = datapipe.fetch_feature(
    feature, node_feature_keys={"author": ["feat"], "paper": ["feat"]}
)
datapipe = datapipe.copy_to(device)
dataloader = gb.DataLoader(datapipe)

训练循环与同构图几乎相同，除了compute_loss的实现，这里将接受两个节点类型和预测的字典。

model = StochasticTwoLayerRGCN(in_features, hidden_features, out_features, etypes)
model = model.to(device)
opt = torch.optim.Adam(model.parameters())

for data in dataloader:
    # For heterogeneous graphs, we need to specify the node types and
    # feature name when accessing the node features. So does the labels.
    input_features = {
        "author": data.node_features[("author", "feat")],
        "paper": data.node_features[("paper", "feat")]
    }
    output_labels = data.labels["paper"]
    output_predictions = model(data.blocks, input_features)
    loss = compute_loss(output_labels, output_predictions)
    opt.zero_grad()
    loss.backward()
    opt.step()

DGL 提供了一个端到端的随机训练示例 RGCN 实现。