GraphormerLayer

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())

基类：Module

具有密集多头注意力的 Graphormer 层，如论文 Do Transformers Really Perform Bad for Graph Representation? 中所介绍。

参数:

• feat_size (int) – 特征尺寸。

• hidden_size (int) – 前馈层的隐藏尺寸。

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

• attn_bias_type (str, 可选) –

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

  • ‘add’ 表示加性注意力偏置。

  • ‘mul’ 表示乘性注意力偏置。

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

• dropout (float, 可选) – Dropout 概率。默认值：0.1。

• attn_dropout (float, 可选) – 注意力 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)

前向计算。

参数:

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

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

• attn_mask (torch.Tensor, 可选) – 用于避免在无效位置上计算的注意力掩码，其中无效位置用 True 值表示。形状：(batch_size, N, N)。注意：对于与不存在节点对应的行，请确保至少有一个条目设置为 False，以防止 softmax 产生 NaNs。

返回:

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

返回类型:

torch.Tensor