dgl.nn.functional.edge_softmax

dgl.nn.functional.edge_softmax(graph, logits, eids='__ALL__', norm_by='dst')[源码]

计算每个节点的入边权重上的 softmax。

对于节点 \(i\)，边 softmax 是一个计算以下值的操作

\[a_{ij} = \frac{\exp(z_{ij})}{\sum_{j\in\mathcal{N}(i)}\exp(z_{ij})}\]

其中 \(z_{ij}\) 是边 \(j\rightarrow i\) 的信号，在 softmax 上下文中也称为 logits。\(\mathcal{N}(i)\) 是指向 \(i\) 的节点的集合。

默认情况下，边 softmax 按目标节点归一化（即在上述公式中，\(ij\) 是节点 i 的入边）。我们也支持按源节点归一化的边 softmax（即在公式中，\(ij\) 是节点 i 的出边）。前一种情况对应于 GAT 和 Transformer 中的 softmax，后一种情况对应于 Capsule 网络中的 softmax。使用边 softmax 的一个示例是图注意力网络 (Graph Attention Network)，其中注意力权重就是用这个操作计算的。使用此操作的其他非 GNN 示例包括Transformer、Capsule 等。

参数:

graph (DGLGraph) – 将执行边 softmax 的图。
logits (torch.Tensor 或 torch.Tensor 字典) – 输入的边特征。异构图可以有一个张量字典，其中每个张量存储相应关系类型的边特征。
eids (torch.Tensor 或 ALL，可选) – 应用边 softmax 的边的 ID。如果是 ALL，则将边 softmax 应用于图中的所有边。默认值：ALL。
norm_by (str, 可以是 src 或 dst) – 按源节点或目标节点归一化。默认值：dst。

返回值:

Softmax 值。

返回类型:

Tensor 或张量元组

注意事项

输入形状: \((E, *, 1)\) 其中 * 表示任意数量的附加维度，\(E\) 等于 eids 的长度。如果 eids 是 ALL，\(E\) 等于图中的边数。
返回形状: \((E, *, 1)\)

同构图示例

以下示例使用 PyTorch 后端。

>>> from dgl.nn.functional import edge_softmax
>>> import dgl
>>> import torch as th

创建一个 DGLGraph 对象并初始化其边特征。

>>> g = dgl.graph((th.tensor([0, 0, 0, 1, 1, 2]), th.tensor([0, 1, 2, 1, 2, 2])))
>>> edata = th.ones(6, 1).float()
>>> edata
    tensor([[1.],
            [1.],
            [1.],
            [1.],
            [1.],
            [1.]])

对图 g 应用边 softmax

>>> edge_softmax(g, edata)
    tensor([[1.0000],
            [0.5000],
            [0.3333],
            [0.5000],
            [0.3333],
            [0.3333]])

对图 g 应用按源节点归一化的边 softmax

>>> edge_softmax(g, edata, norm_by='src')
    tensor([[0.3333],
            [0.3333],
            [0.3333],
            [0.5000],
            [0.5000],
            [1.0000]])

对图 g 的前 4 条边应用边 softmax

>>> edge_softmax(g, edata[:4], th.Tensor([0,1,2,3]))
    tensor([[1.0000],
            [0.5000],
            [1.0000],
            [0.5000]])

异构图示例

创建一个异构图并初始化其边特征。

>>> hg = dgl.heterograph({
...     ('user', 'follows', 'user'): ([0, 0, 1], [0, 1, 2]),
...     ('developer', 'develops', 'game'): ([0, 1], [0, 1])
...     })
>>> edata_follows = th.ones(3, 1).float()
>>> edata_develops = th.ones(2, 1).float()
>>> edata_dict = {('user', 'follows', 'user'): edata_follows,
... ('developer','develops', 'game'): edata_develops}

对图 hg 应用按源节点归一化的边 softmax

>>> edge_softmax(hg, edata_dict, norm_by='src')
    {('developer', 'develops', 'game'): tensor([[1.],
    [1.]]), ('user', 'follows', 'user'): tensor([[0.5000],
    [0.5000],
    [1.0000]])}