mindsponge.cell.TriangleAttention

class mindsponge.cell.TriangleAttention(orientation, num_head, key_dim, gating, layer_norm_dim, batch_size=None, slice_num=0)

三角注意力机制。详细实现过程参考 TriangleAttention 。

氨基酸对ij之间的信息通过ij,ik,jk三条边的信息整合，具体分为投影、自注意力和输出三个步骤，首先进行氨基酸对i,j,k输入的投影，获取i,j,k两两之间的q,k,v，然后通过经典多头自注意机制，在ij氨基酸对之间的信息中添加上i，j，k三角形边之间的关系，最后输出。

参数：

• orientation (int) - 决定三角注意力的方向, 分别作为三角形“starting”和”ending”边的自注意力机制。

• num_head (int) - Attention头的数量。

• key_dim (int) - Attention隐藏层的维度。

• gating (bool) - 判断attention是否经过gating的指示器。

• layer_norm_dim (int) - 归一层的维度。

• batch_size (int) - 三角注意力机制中的batch size参数。默认值： None。

• slice_num (int) - 为了减少内存需要进行切分的数量。默认值： 0。

输入：

• pair_act (Tensor) - pair_act。氨基酸对之间的信息，shape为 \((N_{res}, N_{res}, layer\_norm\_dim)\) 。

• pair_mask (Tensor) - 三角注意力层矩阵的mask。shape为 \((N_{res}, N_{res})\) 。

• index (Tensor) - 在循环中的索引，只会在有控制流的时候使用。

• mask (Tensor) - 当做layernorm操作的时候pair_act的掩码，shape为 \((N_{res}, N_{res})\)，默认值： None。

输出：

Tensor。三角注意力层中的pair_act。shape为 \((N_{res}, N_{res}, layer\_norm\_dim)\) 。

支持平台：

Ascend GPU

样例：

>>> import numpy as np
>>> from mindsponge.cell import TriangleAttention
>>> from mindspore import dtype as mstype
>>> from mindspore import Tensor
>>> model = TriangleAttention(orientation="per_row", num_head=4, key_dim=64, gating=True, layer_norm_dim=64)
>>> input_0 = Tensor(np.ones((256, 256, 64)), mstype.float32)
>>> input_1 = Tensor(np.ones((256, 256)), mstype.float32)
>>> out = model(input_0, input_1, index=0)
>>> print(out.shape)
(256, 256, 64)