7.5 异构图内部原理

本章涵盖了分布式异构图的实现细节。在大多数场景下，这些细节对用户是透明的，但对于高级定制可能很有用。

在 DGL 中，异构图中的节点或边在其各自的节点类型或边类型中具有唯一的 ID。因此，DGL 可以使用一个元组来标识节点或边：(节点/边类型, 类型内 ID)。我们将这种形式的 ID 称为异构 ID。为了对异构图进行分布式训练划分，DGL 会将其转换为同构图，以便我们可以重用为同构图设计的划分算法。每个节点/边因此被唯一映射到一个连续 ID 范围内的整数 ID（例如，从 0 到所有类型节点总数）。我们将转换后的 ID 称为同构 ID。

下面是 ID 转换过程的图示。这里，图有两种类型的节点（\(T0\) 和 \(T1\)），以及四种类型的边（\(R0\)、\(R1\)、\(R2\)、\(R3\)）。图中共有 400 个节点，每种类型有 200 个节点。类型 \(T0\) 的节点 ID 范围是 [0,200)，而类型 \(T1\) 的节点 ID 范围是 [200, 400)。在此示例中，如果我们使用元组来标识节点，类型 \(T0\) 的节点被标识为 (T0, 类型内 ID)，其中类型内 ID 落在 [0, 200) 范围内；类型 \(T1\) 的节点被标识为 (T1, 类型内 ID)，其中类型内 ID 也落在 [0, 200) 范围内。

ID 转换工具

预处理阶段

并行处理管道的所有步骤都使用异构 ID 作为输入和输出。然而，像 ParMETIS 划分这样的一些步骤更容易使用同构 ID 来实现，因此需要一个工具来执行 ID 转换。下面的代码实现了一个简单的 IDConverter，它利用分块图数据格式中元数据 JSON 中的元信息。它从某个节点类型 \(A\) 开始，将其视为节点类型 0，然后为其所有节点分配范围在 \([0, |V_A|-1)\) 的 ID。接着，它移到下一个节点类型 B，将其视为节点类型 1，并为其所有节点分配范围在 \([|V_A|, |V_A|+|V_B|-1)\) 的 ID。

from bisect import bisect_left
import numpy as np

class IDConverter:
    def __init__(self, meta):
        # meta is the JSON object loaded from metadata.json
        self.node_type = meta['node_type']
        self.edge_type = meta['edge_type']
        self.ntype2id_map = {ntype : i for i, ntype in enumerate(self.node_type)}
        self.etype2id_map = {etype : i for i, etype in enumerate(self.edge_type)}
        self.num_nodes = [sum(ns) for ns in meta['num_nodes_per_chunk']]
        self.num_edges = [sum(ns) for ns in meta['num_edges_per_chunk']]
        self.nid_offset = np.cumsum([0] + self.num_nodes)
        self.eid_offset = np.cumsum([0] + self.num_edges)

    def ntype2id(self, ntype):
        """From node type name to node type ID"""
        return self.ntype2id_map[ntype]

    def etype2id(self, etype):
        """From edge type name to edge type ID"""
        return self.etype2id_map[etype]

    def id2ntype(self, id):
        """From node type ID to node type name"""
        return self.node_type[id]

    def id2etype(self, id):
        """From edge type ID to edge type name"""
        return self.edge_type[id]

    def nid_het2hom(self, ntype, id):
        """From heterogeneous node ID to homogeneous node ID"""
        tid = self.ntype2id(ntype)
        if id < 0 or id >= self.num_nodes[tid]:
            raise ValueError(f'Invalid node ID of type {ntype}. Must be within range [0, {self.num_nodes[tid]})')
        return self.nid_offset[tid] + id

    def nid_hom2het(self, id):
        """From heterogeneous node ID to homogeneous node ID"""
        if id < 0 or id >= self.nid_offset[-1]:
            raise ValueError(f'Invalid homogeneous node ID. Must be within range [0, self.nid_offset[-1])')
        tid = bisect_left(self.nid_offset, id) - 1
        # Return a pair (node_type, type_wise_id)
        return self.id2ntype(tid), id - self.nid_offset[tid]

    def eid_het2hom(self, etype, id):
        """From heterogeneous edge ID to homogeneous edge ID"""
        tid = self.etype2id(etype)
        if id < 0 or id >= self.num_edges[tid]:
            raise ValueError(f'Invalid edge ID of type {etype}. Must be within range [0, {self.num_edges[tid]})')
        return self.eid_offset[tid] + id

    def eid_hom2het(self, id):
        """From heterogeneous edge ID to homogeneous edge ID"""
        if id < 0 or id >= self.eid_offset[-1]:
            raise ValueError(f'Invalid homogeneous edge ID. Must be within range [0, self.eid_offset[-1])')
        tid = bisect_left(self.eid_offset, id) - 1
        # Return a pair (edge_type, type_wise_id)
        return self.id2etype(tid), id - self.eid_offset[tid]

加载分区后

将分区加载到训练器或服务器进程后，加载的 GraphPartitionBook 提供了在同构 ID 和异构 ID 之间进行转换的工具。

map_to_per_ntype(): 将同构节点 ID 转换为类型内 ID 和节点类型 ID。
map_to_per_etype(): 将同构边 ID 转换为类型内 ID 和边类型 ID。
map_to_homo_nid(): 将类型内 ID 和节点类型转换为同构节点 ID。
map_to_homo_eid(): 将类型内 ID 和边类型转换为同构边 ID。

由于所有 DGL 的低级分布式图采样操作符都使用同构 ID，DGL 在调用采样操作符之前，会内部将用户指定的异构 ID 转换为同构 ID。下面是一个示例，展示了如何使用 sample_neighbors() 从类型为 "paper" 的节点采样子图。它首先执行 ID 转换，然后在获得采样的子图后，将同构节点/边 ID 转换回异构 ID。

gpb = g.get_partition_book()
# We need to map the type-wise node IDs to homogeneous IDs.
cur = gpb.map_to_homo_nid(seeds, 'paper')
# For a heterogeneous input graph, the returned frontier is stored in
# the homogeneous graph format.
frontier = dgl.distributed.sample_neighbors(g, cur, fanout, replace=False)
block = dgl.to_block(frontier, cur)
cur = block.srcdata[dgl.NID]

block.edata[dgl.EID] = frontier.edata[dgl.EID]
# Map the homogeneous edge Ids to their edge type.
block.edata[dgl.ETYPE], block.edata[dgl.EID] = gpb.map_to_per_etype(block.edata[dgl.EID])
# Map the homogeneous node Ids to their node types and per-type Ids.
block.srcdata[dgl.NTYPE], block.srcdata[dgl.NID] = gpb.map_to_per_ntype(block.srcdata[dgl.NID])
block.dstdata[dgl.NTYPE], block.dstdata[dgl.NID] = gpb.map_to_per_ntype(block.dstdata[dgl.NID])

请注意，从类型 ID 获取节点/边类型很简单 —— 只需从 DistGraph 的 ntypes 属性中获取即可，例如 g.ntypes[node_type_id]。

访问分布式图数据

DistGraph 类支持与 DGLGraph 相似的接口。下面是一个示例，展示了如何获取类型 \(T0\) 的节点 0、10、20 的特征数据。当访问 DistGraph 中的数据时，用户需要使用类型内 ID 以及相应的节点类型或边类型。

import dgl
g = dgl.distributed.DistGraph('graph_name', part_config='data/graph_name.json')
feat = g.nodes['T0'].data['feat'][[0, 10, 20]]

用户可以为特定的节点类型或边类型创建分布式张量和分布式嵌入。分布式张量和嵌入会被分割并存储在多台机器上。要创建一个，用户需要使用 PartitionPolicy 指定其划分方式。默认情况下，DGL 根据第一个维度的大小选择合适的划分策略。但是，如果多个节点类型或边类型具有相同数量的节点或边，DGL 无法自动确定划分策略。用户需要明确指定划分策略。下面展示了一个示例，通过使用 \(T0\) 的划分策略，为节点类型 \(T0\) 创建一个分布式张量，并将其存储为 \(T0\) 的节点数据。

g.nodes['T0'].data['feat1'] = dgl.distributed.DistTensor(
    (g.num_nodes('T0'), 1), th.float32, 'feat1',
    part_policy=g.get_node_partition_policy('T0'))

用于创建分布式张量和嵌入的划分策略是在将异构图加载到图服务器时初始化的。用户不能在运行时创建新的划分策略。因此，用户只能为节点类型或边类型创建分布式张量或嵌入。访问分布式张量和嵌入也需要类型内 ID。