【DGL系列】DGL子图分区的生成和加载API

原创

小锋学长生活大爆炸

发布于 2024-08-17 14:15:11

970

发布于 2024-08-17 14:15:11

文章被收录于专栏：图神经网络

转载请注明出处：小锋学长生活大爆炸[xfxuezhagn.cn] 如果本文帮助到了你，欢迎[点赞、收藏、关注]哦~

7.1 Data Preprocessing — DGL 2.3 documentation

在启动训练作业之前，DGL 要求对输入数据进行分区并分发到目标计算机。为了处理不同尺度的图形，DGL 提供了 2 种分区方法：

用于图形的分区 API，可放入单台计算机内存中。
用于超出单台计算机容量的图形的分布式分区管道。

分区 API

对于相对较小的图，DGL 提供了一个分区 API partition_graph（），用于对内存中的 DGLGraph 对象进行分区。它支持多种分区算法，例如随机分区和 Metis。Metis 分区的好处是，它可以以最小的边缘切割生成分区，从而减少分布式训练和推理的网络通信。DGL 使用最新版本的 Metis，其选项针对具有幂律分布的真实世界图形进行了优化。分区后，API 以易于在训练期间加载的格式构造分区结果。例如

import dgl

g = ...  # create or load a DGLGraph object
dgl.distributed.partition_graph(g, 'mygraph', 2, 'data_root_dir')

将输出以下数据文件：

data_root_dir/
  |-- mygraph.json          # metadata JSON. File name is the given graph name.
  |-- part0/                # data for partition 0
  |  |-- node_feats.dgl     # node features stored in binary format
  |  |-- edge_feats.dgl     # edge features stored in binary format
  |  |-- graph.dgl          # graph structure of this partition stored in binary format
  |
  |-- part1/                # data for partition 1
     |-- node_feats.dgl
     |-- edge_feats.dgl
     |-- graph.dgl

第 7.4 章高级图形分区详细介绍了分区格式。要将分区分发到集群，用户可以将数据保存在所有机器都可以访问的共享文件夹中，或者将元数据 JSON 以及相应的分区文件夹 partX 复制到第 X 台机器。

使用 partition_graph（）需要具有足够大的 CPU RAM 的实例来容纳整个图形结构和特征，这对于具有数千亿条边或大型特征的图形来说可能不可行。接下来，我们将介绍如何使用并行数据准备管道来处理此类情况。

负载均衡

默认情况下，在对图形进行分区时，METIS 仅平衡每个分区中的节点数。这可能会导致配置欠佳，具体取决于手头的任务。例如，在半监督节点分类的情况下，训练器对本地分区中标记节点的子集执行计算。仅平衡图中的节点（有标记和未标记）的分区最终可能会导致计算负载不平衡。为了在每个分区中获得均衡的工作负载，分区 API 允许通过在 partition_graph（）中指定 balance_ntypes，根据每种节点类型的节点数量在分区之间进行均衡。用户可以利用这一点，并考虑训练集、验证集和测试集中是不同的节点类型的情况。

以下示例认为训练集内部的节点和训练集外部的节点是两种类型的节点：

dgl.distributed.partition_graph(g, 'graph_name', 4, '/tmp/test', balance_ntypes=g.ndata['train_mask'])

除了平衡节点类型外，dgl.distributed.partition_graph（）还允许通过指定balance_edges在不同节点类型的节点的度数之间进行平衡。这样可以平衡入射到不同类型节点的边数。

ID 映射

分区后，partition_graph（）会重新映射节点和边 ID，使同一分区的节点排列在一起（在连续的 ID 范围内），从而更容易存储分区的节点/边特征。该 API 还会根据新 ID 自动随机调整节点/边特征。但是，某些下游任务可能希望恢复原始节点/边缘 ID（例如提取计算出的节点嵌入以供以后使用）。对于此类情况，将 return_mapping=True 传递给 partition_graph（），这将使 API 返回重新映射的节点/边 ID 与其原始 ID 之间的 ID 映射。对于同构图，它返回两个向量。第一个向量将每个新节点 ID 映射到其原始 ID;第二个向量将每个新的边 ID 映射到其原始 ID。对于异构图，它返回两个向量字典。第一个字典包含每种节点类型的映射;第二个字典包含每种边类型的映射。

node_map, edge_map = dgl.distributed.partition_graph(g, 'graph_name', 4, '/tmp/test',
                                                     balance_ntypes=g.ndata['train_mask'],
                                                     return_mapping=True)
# Let's assume that node_emb is saved from the distributed training.
orig_node_emb = th.zeros(node_emb.shape, dtype=node_emb.dtype)
orig_node_emb[node_map] = node_emb

加载分区图

DGL 提供了一个 dgl.distributed.load_partition（）函数来加载一个分区进行检查。

import dgl
# load partition 0
part_data = dgl.distributed.load_partition('data_root_dir/graph_name.json', 0)
g, nfeat, efeat, partition_book, graph_name, ntypes, etypes = part_data  # unpack
print(g)

'''
Graph(num_nodes=966043, num_edges=34270118,
      ndata_schemes={'orig_id': Scheme(shape=(), dtype=torch.int64),
                     'part_id': Scheme(shape=(), dtype=torch.int64),
                     '_ID': Scheme(shape=(), dtype=torch.int64),
                     'inner_node': Scheme(shape=(), dtype=torch.int32)}
      edata_schemes={'_ID': Scheme(shape=(), dtype=torch.int64),
                     'inner_edge': Scheme(shape=(), dtype=torch.int8),
                     'orig_id': Scheme(shape=(), dtype=torch.int64)})
'''

如 ID 映射部分所述，每个分区都携带另存为 ndata 或 edata 的辅助信息，例如原始节点/边 ID、分区 ID 等。每个分区不仅保存其拥有的节点/边，还包括与分区相邻的节点/边（称为 HALO 节点/边）。inner_node 和 inner_edge指示节点/边是真正属于分区（值为 True）还是 HALO 节点/边（值为 False）。

load_partition（）函数一次加载所有数据。用户可以分别使用 dgl.distributed.load_partition_feats（）和 dgl.distributed.load_partition_book（） API 加载特征或分区簿。

分布式图分区流水线

暂未用到，先不写了，需要的可以看官方文档

原创声明：本文系作者授权腾讯云开发者社区发表，未经许可，不得转载。

如有侵权，请联系 cloudcommunity@tencent.com 删除。

pytorch

原创声明：本文系作者授权腾讯云开发者社区发表，未经许可，不得转载。

如有侵权，请联系 cloudcommunity@tencent.com 删除。

pytorch

登录后参与评论

0 条评论

热度