GraphBolt技术分享

1.背景

GNN（图神经网络）可以建模图结构数据中的复杂关系和语义信息，它在各种任务中表现出色，包括节点分类、链路预测和图分类。与通常用于图像和视频数据的CNN不同，GNN用于非欧几里德数据，如社交网络。因此，传统的的mini-batch训练方法无法直接应用于GNN。在典型的GNN中，图由结构信息和语义信息组成，图结构通常由稀疏的邻接矩阵表示；语义信息包括节点/边特征，通常表示为向量。如果整个图可以容纳到训练设备的内存中（包括CPU和GPU），则可以通过在每个时期利用完整的图信息来简单地训练模型。然而，现实世界中的图可能非常大，拥有数百万甚至数十亿的节点和边，使得不可能同时将所有信息存储在训练设备的内存中进行训练。为了解决这个问题，研究人员和工程师正在开发新的技术，如图划分、分布式训练和采样策略，以使GNN能够有效处理现实世界的图。

GNN实现Mini-batch的步骤：

• 原始数据准备：准备原始图数据，文件格式多为csv,txt,mat等。
• 数据预处理：将原始数据处理成用于训练的图数据格式，比如DGLgraph。
• mini-batch采样：从整个训练集中随机选择一个子集（即节点、边或图）来计算它们的嵌入。
• 子图采样：基于小批量采样的输出来生成邻居节点，直到达到GNN模型要求的层数。这一步通常不适用于图分类。
• 紧缩：将节点ID紧缩为连续ID序列，并创建紧缩ID与原始ID之间的一对一映射。
• 特征获取：基于子图采样的输出从数据集中获取相关的节点/边特征，一般是输入层的源节点特征。
• 模型训练：给定一个GNN模型，计算节点或边的嵌入，捕捉其在图中的结构和关系特性以及原始特征。根据嵌入进行损失计算，然后进行反向传播和优化。

现存问题：

对于模型训练，DGL和PyG都提供了mini-batch训练的强大支持。然而，从超大规模图中生成batch仍然存在困难，因为可用的存储资源可能会无法容纳大规模图的数据。

2.GraphBolt实现

采样流程：

GraphBolt提供了数据处理和采样的完整datapipe，并且为datapipe中的每个组件提供了明确定义的输入和输出API以及多个标准实现。这使用户可以轻松地将预构建的组件集成到他们的工作流程中，并根据自己的特定需求进行定制。如上图所示，GraphBolt的输入是要训练的完整数据集，具体格式取决于每个组件的实现。输出是一个子数据集，可以用于DGL或PyG的小批量训练。

主要目标：

• 定义datapipe中每个组件的输入和输出API：

1. Dataset :包含完整图结构和语义特征
2. Mini-batch Sampler：基于特定任务的要求，将数据集划分成若干个子集
3. Subgraph Sampler：基于已经划分好的子集，取出相对应的多层子图，用于GNN模型的计算

• 提供了每个组件的规范实现，并且能适应不同的环境：

1. 完整数据集存入内存：适用于小规模图，在保证mini-batch采样顺利进行的前提下，提高采样效率
2. 节点/边特征保存在本地硬盘：适合于大规模图，特征占用存储空间较大，无法一次性全装入内存或显存，允许每次只加载一部分所需的特征进入内存中进行计算。

• 提供每种任务的规范化组件实现

1. 节点分类任务：mini-batch采样主要是对要分类的目标节点进行划分
2. 链接预测任务：mini-batch采样主要针对链接中的正样本和负样本进行生成和采样。
3. 图分类任务：每张图规模相对较小，采样过程类似于PyTorch的采样流程。

主要组件：

• Dataset：访问图数据，图数据可以是内存块或者硬盘上的文件，也可以是远程服务器上的存储单元。Dataset主要由3部分组成：图拓扑，特征存储，splits

1. 图拓扑：保存完整图结构信息，主要是图中的边，后续的子图采样算法是基于完整图拓扑进行的。GraphBolt提供了基于图拓扑的多种不同的采样方法实现，保证了采样过程的性能。
2. 特征存储：以键值对的形式存储节点特征和边特征，基于子图采样的结果，可以在特征存储中获取每个子图需要的输入特征。
3. splits：主要包括训练集、验证集、测试集，用于输入给图神经网络，生成所需要的节点表示和图表示，并且优化图神经网络模型的参数。

• Mini-batch Sampler：以节点分类为例，将要分类的目标节点划分成训练集、验证集、测试集，Mini-batch Sampler的作用是将一个集合中的全部节点划分成若干个子集，每个子集包括若干个节点，这些节点将作为多层图神经网络中最后一层的目标节点，GNN的最终输出也是最后一层目标节点的embedding，用于分类问题的损失计算和模型评估。

• Subgraph Sampler：根据Mini-batch Sampler的采样输出，生成计算最终层目标节点所需要的子图。以节点分类问题为例，使用3层的GNN网络，则需要从最后一层的目标节点开始，采集其一阶邻居形成一张计算依赖图，如此进行3次，共得到3张计算依赖图，分别对应3层GNN单层。值得注意的是，每个计算依赖图都需要把节点ID压缩成从0开始的连续ID编号，并且保存压缩后节点ID和原始图节点ID的一对一映射关系。

• Feature Fetcher：根据子图采样的输出，取出对应的输入特征。以节点分类问题为例，需要提取出多层GNN中第0层的源节点特征作为全部输入特征。同时还要提取最后一层目标节点的标签信息，用于计算损失和评估模型。

• Dataloader：GraphBolt提供的Dataloader实际上是上述4个组件的封装器，将完整图数据集和各类采样器、特征提取器按照计算顺序封装在一起，提供完整的图数据采样和计算的Pipeline。Dataloader的运行过程可以用下图来表示：

3.技术特点

混合CSC格式图拓扑

对于异质图，有两种保存图结构的CSC格式，混合和非混合。非混合格式中，每种边类型都要维持一份自己的CSC实例，同时还有一个映射函数负责将每种边类型映射到自己的CSC实例中。而混合格式中，让所有类型的节点使用同一个连续的ID空间，只需要一个CSC实例就能够保存整张异质图中每种边类型的数据。GraphBolt是采用混合CSC格式来存储图拓扑结构。

内存特征和硬盘特征

GraphBolt提供了两种形式的特征存储，目的是让用户能够尽可能的在一台机器上训练大规模图数据集。内存特征存储会一次性将全部特征先存入内存中，速度相对更快，但是如果图中特征占用空间较大则无法完全放入内存。而硬盘特征存储允许每次只讲计算所需的特征存入内存，这样速度虽然变慢，但是能够很好的处理大规模图数据集，解决了特征无法一次性放入内存的问题。

4.在OpenHGNN上的实践

数据集构造

为测试GraphBolt在OpenHGNN上的性能，可手动将OpenHGNN处理好的DGLGraph转换成GraphBolt所需要的CSC格式。为了方便处理，为异质图中所有节点都补充标签信息。

 for ntype in self.hg.ntypes:
    if ntype != self.cateory:
      self.hg.nodes[ntype].data['label'] = torch.full((self.hg.num_nodes(ntype), ), -1, dtype=torch.int64).to(self.hg.device)
 CSC_Whole_Graph = dgl.graphbolt.from_dglgraph(self.hg)

划分Split

使用GraphBolt提供的ItemDict和ItemSet实例，来构造训练集、验证集、测试集

train_set = gb.ItemSetDict(
{
self.category: gb.ItemSet(
    (train_node_ids_category, train_labels_category),
    names=("seeds", "labels")),
}
)

val_set = gb.ItemSetDict(
{
self.category: gb.ItemSet(
    (val_node_ids_category, val_labels_category),
    names=("seeds", "labels")),
}
)

test_set = gb.ItemSetDict(
{
self.category: gb.ItemSet(
    (test_node_ids_category, test_labels_category),
    names=("seeds", "labels")),
}
)

构造Dataloader

将构造好的数据集（比如train_set）存入datapipe中，经过mini-batch采样（即ItemSampler），子图采样和特征提取，构建完整的mini-batch采样流程，用于为模型提供mini-batch训练数据。

def create_dataloader(itemset, shuffle):    #   itemset可以是train_set,val_set,test_set

    datapipe = gb.ItemSampler(itemset, batch_size=4, shuffle=shuffle)
    datapipe = datapipe.copy_to(self.device, extra_attrs=["seed_nodes"])
    datapipe = datapipe.sample_neighbor(CSC_Whole_Graph, [-1, -1])
    return gb.DataLoader(datapipe)
self.train_dataloader = create_dataloader(train_set, shuffle=True)

模型训练

对Dataloader进行迭代，每次迭代返回一个mini-batch，其中包括了多层GNN模型的每层计算依赖图，最后一层目标节点的标签以及第0层源节点的输入特征，将其输入到GNN模型中，得到输出的节点表示和图表示，用于计算损失和评估模型性能。

 def _mini_train_step(self, ):
  self.model.train()
  loss_all = 0.0
  from tqdm import tqdm
  loader_tqdm = tqdm(self.train_dataloader, ncols=120)

  for i, data in enumerate(loader_tqdm):

      input_nodes = data.input_nodes
      seeds = data.seeds
      if self.to_homo_flag:
          seeds = to_hetero_idx(self.g, self.hg, seeds)
      elif isinstance(input_nodes, dict):
          for key in input_nodes:
              input_nodes[key] = input_nodes[key].to(self.device)
      emb = self.model.input_feature.forward_nodes(input_nodes)
      lbl = self.labels[seeds[self.category]].to(self.device)
      logits = self.model(data.blocks, emb)[self.category]
      loss = self.loss_fn(logits, lbl)
      loss_all += loss.item()
      self.optimizer.zero_grad()
      loss.backward()
      self.optimizer.step()

总结

GraphBolt为图机器学习工作者提供了简洁、易用的数据处理和采样功能实现，使得更大规模的图数据集的处理可以在单机单卡上实现，为图机器学习领域工作者开发利用大规模数据集提供了便利。

实验室算法库简介

两算法库均获得OpenI启智社区“社区优秀孵化项目奖”

GammaGL：

GammaGL是一个基于TensorLayerX实现的开源的多框架的图神经网络（GNN）算法平台，该框架抽象出图数据、消息传递、采样等组件，已经实现了二十余种图数据集和四十余种术界经典GNN算法。GammaGL发布0.4版本 https://github.com/BUPT-GAMMA/GammaGL

OpenHGNN：

OpenHGNN是一个基于深度学习框架PyTorch和图神经网络框架DGL作为底层框架的异质图神经网络（HGNN）工具包。该工具包专注于HGNN的算法模型，抽象出了数据集管理、算法流程、任务评测等模块，提供了超参数调优组件和最优的超参数设置，发布了HGNN设计空间和Benchmark等。OpenHGNN发布0.6版本OpenHGNN：https://github.com/BUPT-GAMMA/OpenHGNN

欢迎 Star ❤ ~