Graphormer层

class dgl.nn.pytorch.gt.GraphormerLayer(feat_size, hidden_size, num_heads, attn_bias_type='add', norm_first=False, dropout=0.1, attn_dropout=0.1, activation=ReLU())

Bases: Module

具有密集多头注意力的Graphormer层，如《Transformer在图表示中真的表现不佳吗？》中介绍的那样。

Parameters:

• feat_size (int) – 特征大小。

• hidden_size (int) – 前馈层的隐藏大小。

• num_heads (int) – 注意力头的数量，feat_size 必须能被其整除。

• attn_bias_type (str, optional) –

  用于修改注意力的注意力偏差类型。可选值为 ‘add’ 或 ‘mul’。默认值为 ‘add’。

  • ’add’ 表示加法注意力偏差。

  • ’mul’ 表示乘法注意力偏差。

• norm_first (bool, optional) – 如果为True，则在注意力和前馈操作之前执行层归一化。否则，它会在之后应用层归一化。默认值：False。

• dropout (float, optional) – 丢弃概率。默认值：0.1。

• attn_dropout (float, optional) – 注意力机制的dropout概率。默认值：0.1。

• activation (可调用的激活层, 可选) – 激活函数。默认值：nn.ReLU()。

示例

>>> import torch as th
>>> from dgl.nn import GraphormerLayer

>>> batch_size = 16
>>> num_nodes = 100
>>> feat_size = 512
>>> num_heads = 8
>>> nfeat = th.rand(batch_size, num_nodes, feat_size)
>>> bias = th.rand(batch_size, num_nodes, num_nodes, num_heads)
>>> net = GraphormerLayer(
        feat_size=feat_size,
        hidden_size=2048,
        num_heads=num_heads
    )
>>> out = net(nfeat, bias)

forward(nfeat, attn_bias=None, attn_mask=None)

前向计算。

Parameters:

• nfeat (torch.Tensor) – 一个3D输入张量。形状：(batch_size, N, feat_size)，其中 N是节点的最大数量。

• attn_bias (torch.Tensor, 可选) – 用于注意力修改的注意力偏置。形状为： (batch_size, N, N, num_heads)。

• attn_mask (torch.Tensor, optional) – 用于避免在无效位置进行计算的注意力掩码，其中无效位置由True值指示。 形状：(batch_size, N, N)。注意：对于对应于不存在节点的行，请确保至少有一个条目设置为False，以防止使用softmax时获得NaN。

Returns:

y – 输出张量。形状：(batch_size, N, feat_size)

Return type:

torch.Tensor