def fuse_add_softmax_dropout()

API接口

def fuse_add_softmax_dropout(training, dropout, attn_mask, attn_scores, attn_head_size, p=0.5, dim=-1):

功能描述

使用NPU自定义算子替换原生写法，以提高性能。

参数说明

• training (Bool)：是否为训练模式。
• dropout (nn.Module): dropout层。
• attn_mask (Tensor)：注意力掩码。
• attn_scores (Tensor)：原始attention分数。
• attn_head_size (Float)：head size。
• p (Float，默认值为0.5)：元素被归零的概率。
• dim (Int，默认值为-1)：待计算softmax的维度。

输出说明

torch.Tensor：mask操作的结果。

示例

调用方式示例：

from torch_npu.contrib.function import fuse_add_softmax_dropout
fuse_add_softmax_dropout(training, dropout, npu_input1, npu_input2, alpha, p=axis)

    >>> training = True
    >>> dropout = nn.DropoutWithByteMask(0.1)
    >>> npu_input1 = torch.rand(96, 12, 384, 384).half().npu()
    >>> npu_input2 = torch.rand(96, 12, 384, 384).half().npu()
    >>> alpha = 0.125
    >>> axis = -1
    >>> output = fuse_add_softmax_dropout(training, dropout, npu_input1, npu_input2, alpha, p=axis)