dgl.nn.functional.edge_softmax

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

计算每个节点的传入边权重的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等。

Parameters:

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

Returns:

Softmax 值。

Return type:

张量或元组的张量

注释

输入形状：\((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]])}