mindsponge.cell.GlobalAttention

class mindsponge.cell.GlobalAttention(num_head, gating, input_dim, output_dim, batch_size=None)

global gated自注意力机制，具体实现请参考 Highly accurate protein structure prediction with AlphaFold 。对于GlobalAttention模块，query/key/value tensor的shape需保持一致。

参数：

• num_head (int) - 头的数量。

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

• input_dim (int) - 输入的最后一维的长度。

• output_dim (int) - 输出的最后一维的长度。

• batch_size (int) - attention中权重的batch size，仅在有while控制流时使用，默认值： None。

输入：

• q_data (Tensor) - shape为 \((batch\_size, seq\_length, input\_dim)\) 的query tensor，其中seq_length是query向量的序列长度。

• m_data (Tensor) - shape为 \((batch\_size, seq\_length, input\_dim)\) 的key和value tensor。

• q_mask (Tensor) - shape为 \((batch\_size, seq\_length, 1)\) 的q_data的mask。

• bias (Tensor) - attention矩阵的偏置。默认值： None。

• index (Tensor) - 在while循环中的索引，仅在有while控制流时使用。默认值： None。

输出：

Tensor。GlobalAttention层的输出tensor，shape是 \((batch\_size, seq\_length, output\_dim)\)。

支持平台：

Ascend GPU

样例：

>>> import numpy as np
>>> from mindsponge.cell import GlobalAttention
>>> from mindspore import dtype as mstype
>>> from mindspore import Tensor
>>> model = GlobalAttention(num_head=4, input_dim=64, gating=True, output_dim=256)
>>> q_data = Tensor(np.ones((32, 128, 64)), mstype.float32)
>>> m_data = Tensor(np.ones((32, 128, 64)), mstype.float32)
>>> q_mask = Tensor(np.ones((32, 128, 1)), mstype.float32)
>>> attn_out= model(q_data, m_data, q_mask)
>>> print(attn_out.shape)
(32, 128, 256)