7.1 分布式训练所需的图数据预处理

(English Version)

DGL要求预处理图数据以进行分布式训练,这包括两个步骤:1)将一张图划分为多张子图(分区),2)为节点和边分配新的ID。 DGL提供了一个API以执行这两个步骤。该API支持随机划分和一个基于 Metis 的划分。Metis划分的好处在于, 它可以用最少的边分割以生成分区,从而减少了用于分布式训练和推理的网络通信。DGL使用最新版本的Metis, 并针对真实世界中具有幂律分布的图进行了优化。在图划分后,API以易于在训练期间加载的格式构造划分结果。

Note: 图划分API当前在一台机器上运行。 因此如果一张图很大,用户将需要一台大内存的机器来对图进行划分。 如果图大到不能在一台机器上完成分图,DGL支持分布式图划分,目前基于ParMETIS,具体步骤参考 https://doc.dgl.ai/guide/distributed-preprocessing.html#distributed-partitioning

默认情况下,为了在分布式训练/推理期间定位节点/边,API将新ID分配给输入图的节点和边。 分配ID后,该API会相应地打乱所有节点数据和边数据。在训练期间,用户只需使用新的节点和边的ID。 与此同时,用户仍然可以通过 g.ndata['orig_id']g.edata['orig_id'] 获取原始ID。 其中 gDistGraph 对象(详细解释,请参见:ref:guide-distributed-apis)。

DGL将图划分结果存储在输出目录中的多个文件中。输出目录里始终包含一个名为xxx.json的JSON文件,其中xxx是提供给划分API的图的名称。 JSON文件包含所有划分的配置。如果该API没有为节点和边分配新ID,它将生成两个额外的NumPy文件:node_map.npyedge_map.npy。 它们存储节点和边ID与分区ID之间的映射。对于具有十亿级数量节点和边的图,两个文件中的NumPy数组会很大, 这是因为图中的每个节点和边都对应一个条目。在每个分区的文件夹内,有3个文件以DGL格式存储分区数据。 graph.dgl 存储分区的图结构以及节点和边上的一些元数据。node_feats.dgledge_feats.dgl 存储属于该分区的节点和边的所有特征。

data_root_dir/
    |-- xxx.json             # JSON中的分区配置文件
    |-- node_map.npy         # 存储在NumPy数组中的每个节点的分区ID(可选)
    |-- edge_map.npy         # 存储在NumPy数组中的每个边的分区ID(可选)
    |-- part0/               # 分区0的数据
        |-- node_feats.dgl   # 以二进制格式存储的节点特征
        |-- edge_feats.dgl   # 以二进制格式存储的边特征
        |-- graph.dgl        # 以二进制格式存储的子图结构
    |-- part1/               # 分区1的数据
        |-- node_feats.dgl
        |-- edge_feats.dgl
        |-- graph.dgl

负载均衡

在对图进行划分时,默认情况下,Metis仅平衡每个子图中的节点数。根据当前的任务情况,这可能带来非最优的配置。 例如,在半监督节点分类的场景里,训练器会对局部分区中带标签节点的子集进行计算。 一个仅平衡图中节点(带标签和未带标签)的划分可能会导致计算负载不平衡。为了在每个分区中获得平衡的工作负载, 划分API通过在 dgl.distributed.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 来平衡每个类型节点在子图中的入度。这平衡了不同类型节点的连边数量。

Note: 传给 dgl.distributed.partition_graph() 的图名称是一个重要的参数。 dgl.distributed.DistGraph 使用该名称来识别一个分布式的图。一个有效的图名称应该仅包含字母和下划线。