PyTorch Geometric - 如何对巨大的图进行采样以使用小批量训练 GNN

Question

我想使用 PyTorch Geometric 在一个巨大的图（大约 1M 个节点）上进行节点回归，但是我无法创建一个

Data

对象，因为完整的图不在 RAM 中，所以我不能将

DataLoader

类用于 mini -批处理和培训。

一些例子（例如4. Scaling Graph Neural Networks）介绍了

Clusterdata

和

ClusterLoader

类，但这对我的情况没有帮助，因为它们实际上可以加载整个图。

我已经将节点嵌入和边预先计算到单独的文件中，我可以非常快速地读取这些文件，以加载图形子集和特定节点的嵌入。但是，我不知道我应该如何在训练期间对图表进行采样，或者是否有任何现有的 PyTorch 模块已经这样做了。

我的问题是：PyTorch Geometric 是否有任何模块可以创建小批量来训练我的 GCN 而无需将整个图形加载到内存中？如果不是，我应该如何进行图表抽样？

在 PyTorch Geometric 文档中有许多节点和图形回归、分类的示例......但是没有一个示例解释如何处理如此大的图形，因为它们使用由许多小图形组成的数据集，这些图形都适合 RAM。

在另一个 Google Colab notebook 示例中（2. 使用图神经网络进行节点分类），现有数据集 (Planetoid) 中的整个图被加载到 RAM 中。

dataset = Planetoid(root='data/Planetoid', name='Cora', transform=NormalizeFeatures())
...
data = dataset[0]  # Get the first graph object.

然后，为模型的一个训练时期定义一个

train

函数，它使用图形的完整

data.x

和

data.edge_index

。

def train():
      model.train()
      optimizer.zero_grad()  # Clear gradients.
      out = model(data.x, data.edge_index)  # Perform a single forward pass.
      loss = criterion(out[data.train_mask], data.y[data.train_mask])  # Compute the loss solely based on the training nodes.
      loss.backward()  # Derive gradients.
      optimizer.step()  # Update parameters based on gradients.
      return loss

从这个例子中，我猜我图中的每个小批量应该由以下人员创建：

1. 从我的图中选择节点的随机子集。
2. 读取涉及所选节点的所有边和节点嵌入。

然后，在这个 mini-batch 上训练一个 epoch。但是，如果我随机选择的节点都不相邻，并且没有完成消息传递怎么办？有没有正确的方法从这个图中采样而不发生这种情况？例如，我们可以选择一个随机节点，然后将某个邻域作为子集吗？