mindsponge.cell.InvariantPointAttention

查看源文件
class mindsponge.cell.InvariantPointAttention(num_head, num_scalar_qk, num_scalar_v, num_point_v, num_point_qk, num_channel, pair_dim)[源代码]

该模块用于更新序列表示(即输入inputs_1d),在序列表示中加入位置信息。 其中注意力由三部分构成,即由序列表示得到的q, k, v,由序列表示与刚体变换组局部坐标系T交互得到的q’, k’, v’, 以及从氨基酸对表示(输入中的inputs_2d)中得到的偏移b。

\[a_{ij} = Softmax(w_l(c_1{q_i}^Tk_j+b{ij}-c_2\sum {\left \| T_i\circ q'_i-T_j\circ k'_j \right \| ^{2 } }))\]

其中i,j分别表示序列中第i、第j个氨基酸,T即输入中的rotation和translation。

参考文献:Jumper et al. (2021) Suppl. Alg. 22 InvariantPointAttention

参数:
  • num_head (int) - 头的数量。

  • num_scalar_qk (int) - scalar query/key的数量。

  • num_scalar_v (int) - scalar value的数量。

  • num_point_v (int) - point value的数量。

  • num_point_qk (int) - point query/key的数量。

  • num_channel (int) - 通道数量。

  • pair_dim (int) - pair的最后一维长度。

输入:
  • inputs_1d (Tensor) - Evoformer模块的输出msa表示矩阵中的第一行,也即序列表示, \([N_{res}, num\_channel]\)

  • inputs_2d (Tensor) - Evoformor模块的输出氨基酸对表示矩阵, \([N_{res}, N_{res}, pair\_dim]\)

  • mask (Tensor) - 掩码,表示inputs_1d的哪些元素参与了attention, \([N_{res}, 1]\)

  • rotation (tuple) - 刚体变换组局部坐标系 \(T(r,t)\) 中的旋转信息, 长度为9的元组,每个元素shape为 \([N_{res}]\)

  • translation (tuple) - 刚体变换组局部坐标系 \(T(r,t)\) 中的旋转信息的偏移信息, 长度为3的元组,每个元素shape为 \([N_{res}]\)

输出:

Tensor,input_1d的更新,shape为 \([N_{res}, num\_channel]\)

支持平台:

Ascend GPU

样例:

>>> import numpy as np
>>> from mindsponge.cell import InvariantPointAttention
>>> from mindspore import dtype as mstype
>>> from mindspore import Tensor
>>> import mindspore.context as context
>>> context.set_context(mode=context.GRAPH_MODE)
>>> model = InvariantPointAttention(num_head=12, num_scalar_qk=16, num_scalar_v=16,
...                                 num_point_v=8, num_point_qk=4,
...                                 num_channel=384, pair_dim=128)
>>> inputs_1d = Tensor(np.ones((256, 384)), mstype.float32)
>>> inputs_2d = Tensor(np.ones((256, 256, 128)), mstype.float32)
>>> mask = Tensor(np.ones((256, 1)), mstype.float32)
>>> rotation = tuple([Tensor(np.ones(256), mstype.float16) for _ in range(9)])
>>> translation = tuple([Tensor(np.ones(256), mstype.float16) for _ in range(3)])
>>> attn_out = model(inputs_1d, inputs_2d, mask, rotation, translation)
>>> print(attn_out.shape)
(256, 384)