[240905] 如何使用 JAX 和 Equinox 构建图卷积网络 | Cascadia 字体家族迎来新成员

x-cmd

于 2024-09-05 08:45:17 发布

阅读量551

点赞数 5

分类专栏： daily blog 文章标签：网络邻接矩阵边列表 JIT 微软字体开源

本文链接：https://blog.csdn.net/edwinjhlee/article/details/141916303

版权

daily blog 专栏收录该内容

365 篇文章

订阅专栏

如何使用 JAX 和 Equinox 构建图卷积网络

文章介绍如何使用 JAX 和 Equinox 构建图卷积网络 (GNN)。我们将分别使用邻接矩阵和边列表两种方式实现图卷积层，并比较它们的优缺点。

1 使用邻接矩阵

1.1 邻接矩阵表示法

对于包含 NNN 个节点的图，我们可以使用一个 N×NN \times NN×N 的邻接矩阵 AAA 来表示节点之间的关系。矩阵元素 ai,j∈{0,1}a_{i, j} \in \{0, 1\}ai,j\u200b∈{0,1} 表示从节点 jjj 到节点 iii 是否存在边 (1 表示存在，0 表示不存在)。

1.2 图卷积层实现

import jax.experimental.sparse as jsparse

class GraphConv(eqx.Module):
    linear: nn.Linear

    def __init__(self, hidden_dim: int, *, key: PRNGKeyArray):
        self.linear = nn.Linear(hidden_dim, hidden_dim, key=key)

    def __call__(
        self,
        nodes: Float[Array, "n_nodes hidden_dim"],
        adjacency: Int[jsparse.BCOO, "n_nodes n_nodes"]
    ) -> Float[Array, "n_nodes hidden_dim"]:
        messages = vmap(self.linear)(nodes)
        return adjacency @ messages

1.3 计算过程解释

上述代码中，我们首先对所有节点应用线性变换，然后将结果与邻接矩阵相乘。这等效于对每个节点 iii，将其所有邻居节点 jjj 的特征进行加权求和，其中权重由线性变换矩阵 WWW 决定。

2 使用边列表

2.1 边列表表示法

边列表使用一个 M×2M \times 2M×2 的张量 EEE 来表示图中的边。其中， ek=(j,i)e_k = (j, i)ek\u200b=(j,i) 表示第 kkk 条边是从节点 jjj 到节点 iii。

2.2 图卷积层实现

class GraphConv(eqx.Module):
    linear: nn.Linear

    def __init__(self, hidden_dim: int, *, key: PRNGKeyArray):
        self.linear = nn.Linear(hidden_dim, hidden_dim, key=key)

    def __call__(
        self,
        nodes: Float[Array, "n_nodes hidden_dim"],
        edges: Int[Array, "n_edges 2"],
    ) -> Float[Array, "n_nodes hidden_dim"]:
        messages = vmap(self.linear)(nodes)
        messages = messages[edges[:, 0]]  # 获取源节点特征
        messages = jax.ops.segment_sum(
            data=messages,
            segment_ids=edges[:, 1],
            num_segments=len(nodes),
        )  # 按目标节点聚合特征
        return messages

2.3 代码解析：`jax.ops.segment_sum`

jax.ops.segment_sum 函数用于根据 segment_ids 对数据进行分组求和。在本例中，我们将所有边的源节点特征按照目标节点 ID 进行分组求和，从而得到每个目标节点的聚合特征。

2.4 计算节点度数示例

ones = jnp.ones(len(edges), dtype=jnp.int32)
degrees = jax.ops.segment_sum(
    data=ones,
    segment_ids=edges[:, 1],
    num_segments=len(nodes),
)