| import math | |
| import torch | |
| import torch.nn as nn | |
| # from torch.nn import Module | |
| # # for gzlabel contable_gpu env | |
| # class MultiheadAttention(Module): | |
| # r"""Allows the model to jointly attend to information | |
| # from different representation subspaces. | |
| # See `Attention Is All You Need <https://arxiv.org/abs/1706.03762>`_. | |
| # | |
| # .. math:: | |
| # \text{MultiHead}(Q, K, V) = \text{Concat}(head_1,\dots,head_h)W^O | |
| # | |
| # where :math:`head_i = \text{Attention}(QW_i^Q, KW_i^K, VW_i^V)`. | |
| # | |
| # Args: | |
| # embed_dim: Total dimension of the model. | |
| # num_heads: Number of parallel attention heads. Note that ``embed_dim`` will be split | |
| # across ``num_heads`` (i.e. each head will have dimension ``embed_dim // num_heads``). | |
| # dropout: Dropout probability on ``attn_output_weights``. Default: ``0.0`` (no dropout). | |
| # bias: If specified, adds bias to input / output projection layers. Default: ``True``. | |
| # add_bias_kv: If specified, adds bias to the key and value sequences at dim=0. Default: ``False``. | |
| # add_zero_attn: If specified, adds a new batch of zeros to the key and value sequences at dim=1. | |
| # Default: ``False``. | |
| # kdim: Total number of features for keys. Default: ``None`` (uses ``kdim=embed_dim``). | |
| # vdim: Total number of features for values. Default: ``None`` (uses ``vdim=embed_dim``). | |
| # batch_first: If ``True``, then the input and output tensors are provided | |
| # as (batch, seq, feature). Default: ``False`` (seq, batch, feature). | |
| # | |
| # Examples:: | |
| # | |
| # >>> multihead_attn = nn.MultiheadAttention(embed_dim, num_heads) | |
| # >>> attn_output, attn_output_weights = multihead_attn(query, key, value) | |
| # """ | |
| # __constants__ = ['batch_first'] | |
| # bias_k: Optional[torch.Tensor] | |
| # bias_v: Optional[torch.Tensor] | |
| # | |
| # def __init__(self, embed_dim, num_heads, dropout=0., bias=True, add_bias_kv=False, add_zero_attn=False, | |
| # kdim=None, vdim=None, batch_first=False, device=None, dtype=None) -> None: | |
| # factory_kwargs = {'device': device, 'dtype': dtype} | |
| # super(MultiheadAttention, self).__init__() | |
| # self.embed_dim = embed_dim | |
| # self.kdim = kdim if kdim is not None else embed_dim | |
| # self.vdim = vdim if vdim is not None else embed_dim | |
| # self._qkv_same_embed_dim = self.kdim == embed_dim and self.vdim == embed_dim | |
| # | |
| # self.num_heads = num_heads | |
| # self.dropout = dropout | |
| # self.batch_first = batch_first | |
| # self.head_dim = embed_dim // num_heads | |
| # assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads" | |
| # | |
| # if self._qkv_same_embed_dim is False: | |
| # self.q_proj_weight = Parameter(torch.empty((embed_dim, embed_dim), **factory_kwargs)) | |
| # self.k_proj_weight = Parameter(torch.empty((embed_dim, self.kdim), **factory_kwargs)) | |
| # self.v_proj_weight = Parameter(torch.empty((embed_dim, self.vdim), **factory_kwargs)) | |
| # self.register_parameter('in_proj_weight', None) | |
| # else: | |
| # self.in_proj_weight = Parameter(torch.empty((3 * embed_dim, embed_dim), **factory_kwargs)) | |
| # self.register_parameter('q_proj_weight', None) | |
| # self.register_parameter('k_proj_weight', None) | |
| # self.register_parameter('v_proj_weight', None) | |
| # | |
| # if bias: | |
| # self.in_proj_bias = Parameter(torch.empty(3 * embed_dim, **factory_kwargs)) | |
| # else: | |
| # self.register_parameter('in_proj_bias', None) | |
| # self.out_proj = NonDynamicallyQuantizableLinear(embed_dim, embed_dim, bias=bias, **factory_kwargs) | |
| # | |
| # if add_bias_kv: | |
| # self.bias_k = Parameter(torch.empty((1, 1, embed_dim), **factory_kwargs)) | |
| # self.bias_v = Parameter(torch.empty((1, 1, embed_dim), **factory_kwargs)) | |
| # else: | |
| # self.bias_k = self.bias_v = None | |
| # | |
| # self.add_zero_attn = add_zero_attn | |
| # | |
| # self._reset_parameters() | |
| # | |
| # def _reset_parameters(self): | |
| # if self._qkv_same_embed_dim: | |
| # xavier_uniform_(self.in_proj_weight) | |
| # else: | |
| # xavier_uniform_(self.q_proj_weight) | |
| # xavier_uniform_(self.k_proj_weight) | |
| # xavier_uniform_(self.v_proj_weight) | |
| # | |
| # if self.in_proj_bias is not None: | |
| # constant_(self.in_proj_bias, 0.) | |
| # constant_(self.out_proj.bias, 0.) | |
| # if self.bias_k is not None: | |
| # xavier_normal_(self.bias_k) | |
| # if self.bias_v is not None: | |
| # xavier_normal_(self.bias_v) | |
| # | |
| # def __setstate__(self, state): | |
| # # Support loading old MultiheadAttention checkpoints generated by v1.1.0 | |
| # if '_qkv_same_embed_dim' not in state: | |
| # state['_qkv_same_embed_dim'] = True | |
| # | |
| # super(MultiheadAttention, self).__setstate__(state) | |
| # | |
| # def forward(self, query: Tensor, key: Tensor, value: Tensor, key_padding_mask: Optional[Tensor] = None, | |
| # need_weights: bool = True, attn_mask: Optional[Tensor] = None) -> Tuple[Tensor, Optional[Tensor]]: | |
| # r""" | |
| # Args: | |
| # query: Query embeddings of shape :math:`(L, N, E_q)` when ``batch_first=False`` or :math:`(N, L, E_q)` | |
| # when ``batch_first=True``, where :math:`L` is the target sequence length, :math:`N` is the batch size, | |
| # and :math:`E_q` is the query embedding dimension ``embed_dim``. Queries are compared against | |
| # key-value pairs to produce the output. See "Attention Is All You Need" for more details. | |
| # key: Key embeddings of shape :math:`(S, N, E_k)` when ``batch_first=False`` or :math:`(N, S, E_k)` when | |
| # ``batch_first=True``, where :math:`S` is the source sequence length, :math:`N` is the batch size, and | |
| # :math:`E_k` is the key embedding dimension ``kdim``. See "Attention Is All You Need" for more details. | |
| # value: Value embeddings of shape :math:`(S, N, E_v)` when ``batch_first=False`` or :math:`(N, S, E_v)` when | |
| # ``batch_first=True``, where :math:`S` is the source sequence length, :math:`N` is the batch size, and | |
| # :math:`E_v` is the value embedding dimension ``vdim``. See "Attention Is All You Need" for more details. | |
| # key_padding_mask: If specified, a mask of shape :math:`(N, S)` indicating which elements within ``key`` | |
| # to ignore for the purpose of attention (i.e. treat as "padding"). Binary and byte masks are supported. | |
| # For a binary mask, a ``True`` value indicates that the corresponding ``key`` value will be ignored for | |
| # the purpose of attention. For a byte mask, a non-zero value indicates that the corresponding ``key`` | |
| # value will be ignored. | |
| # need_weights: If specified, returns ``attn_output_weights`` in addition to ``attn_outputs``. | |
| # Default: ``True``. | |
| # attn_mask: If specified, a 2D or 3D mask preventing attention to certain positions. Must be of shape | |
| # :math:`(L, S)` or :math:`(N\cdot\text{num\_heads}, L, S)`, where :math:`N` is the batch size, | |
| # :math:`L` is the target sequence length, and :math:`S` is the source sequence length. A 2D mask will be | |
| # broadcasted across the batch while a 3D mask allows for a different mask for each entry in the batch. | |
| # Binary, byte, and float masks are supported. For a binary mask, a ``True`` value indicates that the | |
| # corresponding position is not allowed to attend. For a byte mask, a non-zero value indicates that the | |
| # corresponding position is not allowed to attend. For a float mask, the mask values will be added to | |
| # the attention weight. | |
| # | |
| # Outputs: | |
| # - **attn_output** - Attention outputs of shape :math:`(L, N, E)` when ``batch_first=False`` or | |
| # :math:`(N, L, E)` when ``batch_first=True``, where :math:`L` is the target sequence length, :math:`N` is | |
| # the batch size, and :math:`E` is the embedding dimension ``embed_dim``. | |
| # - **attn_output_weights** - Attention output weights of shape :math:`(N, L, S)`, where :math:`N` is the batch | |
| # size, :math:`L` is the target sequence length, and :math:`S` is the source sequence length. Only returned | |
| # when ``need_weights=True``. | |
| # """ | |
| # if self.batch_first: | |
| # query, key, value = [x.transpose(1, 0) for x in (query, key, value)] | |
| # | |
| # if not self._qkv_same_embed_dim: | |
| # attn_output, attn_output_weights = F.multi_head_attention_forward( | |
| # query, key, value, self.embed_dim, self.num_heads, | |
| # self.in_proj_weight, self.in_proj_bias, | |
| # self.bias_k, self.bias_v, self.add_zero_attn, | |
| # self.dropout, self.out_proj.weight, self.out_proj.bias, | |
| # training=self.training, | |
| # key_padding_mask=key_padding_mask, need_weights=need_weights, | |
| # attn_mask=attn_mask, use_separate_proj_weight=True, | |
| # q_proj_weight=self.q_proj_weight, k_proj_weight=self.k_proj_weight, | |
| # v_proj_weight=self.v_proj_weight) | |
| # else: | |
| # attn_output, attn_output_weights = F.multi_head_attention_forward( | |
| # query, key, value, self.embed_dim, self.num_heads, | |
| # self.in_proj_weight, self.in_proj_bias, | |
| # self.bias_k, self.bias_v, self.add_zero_attn, | |
| # self.dropout, self.out_proj.weight, self.out_proj.bias, | |
| # training=self.training, | |
| # key_padding_mask=key_padding_mask, need_weights=need_weights, | |
| # attn_mask=attn_mask) | |
| # if self.batch_first: | |
| # return attn_output.transpose(1, 0), attn_output_weights | |
| # else: | |
| # return attn_output, attn_output_weights | |
| class PositionalEncoding(nn.Module): | |
| "Implement the PE function." | |
| def __init__(self, d_model, dropout, max_len=5000): | |
| #d_model=512,dropout=0.1, | |
| #max_len=5000代表事先准备好长度为5000的序列的位置编码，其实没必要， | |
| #一般100或者200足够了。 | |
| super(PositionalEncoding, self).__init__() | |
| self.dropout = nn.Dropout(p=dropout) | |
| # Compute the positional encodings once in log space. | |
| pe = torch.zeros(max_len, d_model) | |
| #(5000,512)矩阵，保持每个位置的位置编码，一共5000个位置， | |
| #每个位置用一个512维度向量来表示其位置编码 | |
| position = torch.arange(0, max_len).unsqueeze(1) | |
| # (5000) -> (5000,1) | |
| div_term = torch.exp(torch.arange(0, d_model, 2) * | |
| -(math.log(10000.0) / d_model)) | |
| # (0,2,…, 4998)一共准备2500个值，供sin, cos调用 | |
| pe[:, 0::2] = torch.sin(position * div_term) # 偶数下标的位置 | |
| pe[:, 1::2] = torch.cos(position * div_term) # 奇数下标的位置 | |
| pe = pe.unsqueeze(0) | |
| # (5000, 512) -> (1, 5000, 512) 为batch.size留出位置 | |
| self.register_buffer('pe', pe) | |
| def forward(self, x): | |
| x = x + self.pe[:, :x.size(1)] | |
| # 接受1.Embeddings的词嵌入结果x， | |
| #然后把自己的位置编码pe，封装成torch的Variable(不需要梯度)，加上去。 | |
| #例如，假设x是(30,10,512)的一个tensor， | |
| #30是batch.size, 10是该batch的序列长度, 512是每个词的词嵌入向量； | |
| #则该行代码的第二项是(1, min(10, 5000), 512)=(1,10,512)， | |
| #在具体相加的时候，会扩展(1,10,512)为(30,10,512)， | |
| #保证一个batch中的30个序列，都使用（叠加）一样的位置编码。 | |
| return self.dropout(x) # 增加一次dropout操作 | |
| # 注意，位置编码不会更新，是写死的，所以这个class里面没有可训练的参数。 | |
| class TwoTrackAttention(nn.Module): | |
| def __init__(self, d_attn, n_head, d_ff=512, dropout=0.1) -> None: | |
| super().__init__() | |
| self.self_attn = torch.nn.MultiheadAttention( | |
| d_attn, n_head, | |
| dropout = dropout, | |
| batch_first=True # gzbl 这边的pytorch版本没有这个参数 | |
| ) | |
| self.dropout_self = nn.Dropout(dropout) | |
| self.cross_attn = torch.nn.MultiheadAttention( | |
| d_attn, n_head, | |
| dropout = dropout, | |
| batch_first=True | |
| ) | |
| self.dropout_cross = nn.Dropout(dropout) | |
| self.norm1 = nn.LayerNorm(d_attn) | |
| self.ff1 = nn.Linear(d_attn, d_ff) | |
| self.dropout_ff = nn.Dropout(dropout) | |
| self.ff2 = nn.Linear(d_ff, d_attn) | |
| self.norm2 = nn.LayerNorm(d_attn) | |
| self.dropout = nn.Dropout(dropout) | |
| self.activation = nn.ReLU() | |
| # self.s_query = nn.Linear(d_attn,d_attn) | |
| # self.s_key = nn.Linear(d_attn,d_attn) | |
| # self.s_value = nn.Linear(d_attn,d_attn) | |
| # | |
| # self.c_query = nn.Linear(d_attn,d_attn) | |
| # self.c_key = nn.Linear(d_attn,d_attn) | |
| # self.c_value = nn.Linear(d_attn,d_attn) | |
| def forward(self, obj_update, obj_message): | |
| self_update = self.self_attn( | |
| query = obj_update, | |
| key = obj_update, | |
| value = obj_update | |
| )[0] | |
| cross_update = self.cross_attn( | |
| query = obj_update, # [1, 299, 128] | |
| key = obj_message, # [1, 74, 128] | |
| value = obj_message # [1, 74, 128] | |
| )[0] | |
| # [torch.Size([1, 299, 128]), torch.Size([1, 74, 128]), torch.Size([1, 74, 128])] | |
| obj_update = obj_update + self.dropout_self(self_update) + self.dropout_cross(cross_update) | |
| obj_update = self.norm1(obj_update) | |
| ff_update = self.ff2(self.dropout_ff(self.activation(self.ff1(obj_update)))) | |
| obj_update = obj_update + self.dropout(ff_update) | |
| obj_update = self.norm2(obj_update) | |
| return obj_update | |
| class SymertricTwoTrackAttention(nn.Module): | |
| def __init__(self, d_attn, n_head, d_ff=512, dropout=0.1,sync = False) -> None: | |
| super().__init__() | |
| self.tta1 = TwoTrackAttention(d_attn, n_head, d_ff, dropout) | |
| self.tta2 = TwoTrackAttention(d_attn, n_head, d_ff, dropout) | |
| self.sync = sync | |
| def forward(self, obj_1, obj_2): | |
| if self.sync: | |
| return self.tta1(obj_1, obj_2), self.tta2(obj_2, obj_1) | |
| else: | |
| obj_1 = self.tta1(obj_1, obj_2) | |
| obj_2 = self.tta2(obj_2, obj_1) | |
| return obj_1, obj_2 | |
| class LinearFF(nn.Module): | |
| def __init__(self, d_in, d_out, dropout=0.1) -> None: | |
| super().__init__() | |
| self.emb = nn.Linear(d_in, d_out) | |
| self.norm = nn.LayerNorm(d_out) | |
| self.dropout = nn.Dropout(dropout) | |
| self.activation = nn.ReLU() | |
| def forward(self, f_in): | |
| f_in = f_in.permute(0,2,1) | |
| return self.norm(self.dropout(self.activation(self.emb(f_in)))) | |
| class ProteinRNAInteraction(nn.Module): | |
| def __init__(self, d_pro, d_rna, n_layers, d_attn, n_head=4, d_ff=512, dropout=0.1,sync=False) -> None: | |
| super().__init__() | |
| print('sync update ProteinRNAInteraction',sync) | |
| self.pro_emb = LinearFF(d_pro, d_attn) | |
| self.pro_rna = LinearFF(d_rna, d_attn) | |
| self.pro_pos = PositionalEncoding(d_attn,dropout) | |
| self.rna_pos = PositionalEncoding(d_attn,dropout) | |
| self.layers = nn.ModuleList([ | |
| SymertricTwoTrackAttention(d_attn, n_head, d_ff, dropout,sync = sync) for _ in range(n_layers) | |
| ]) | |
| self.pred = nn.Linear(d_attn, 1) | |
| # self.pred = nn.Linear(2*d_attn, 1) | |
| self.sigmoid = nn.Sigmoid() | |
| def forward(self, f_pro, f_rna): | |
| # print(f_pro.shape) | |
| # print(f_pro.device) | |
| f_pro = self.pro_emb(f_pro) | |
| f_rna = self.pro_rna(f_rna) | |
| f_pro = self.pro_pos(f_pro) | |
| f_rna = self.rna_pos(f_rna) | |
| for layer in self.layers: | |
| f_pro, f_rna = layer(f_pro, f_rna) | |
| f_pro = f_pro.unsqueeze(2) # [B, L, R, D] | |
| f_rna = f_rna.unsqueeze(1) | |
| prob = self.sigmoid(self.pred(f_rna.mul(f_pro))) | |
| return prob | |
| # f_pro = f_pro.unsqueeze(2) # [1, 299, 1, 128] | |
| # f_rna = f_rna.unsqueeze(1) # [1, 1, 74, 128] | |
| # f_pro = f_pro.repeat(1, 1, f_rna.shape[2], 1) # [B, L, R, D] | |
| # f_rna = f_rna.repeat(1, f_pro.shape[1], 1, 1) # [B, L, R, D] | |
| # | |
| # # prob = self.pred(f_rna.mul(f_pro)) | |
| # prob = self.pred(torch.cat([f_pro, f_rna], -1)) | |
| # # print(prob.max(),prob.min(),prob.mean()) | |
| # prob = torch.sigmoid(prob) | |
| # # prob = self.sigmoid(prob) | |
| # # prob = self.sigmoid(self.pred(torch.cat([f_pro, f_rna], -1))) # pred : -0.06, 0.619 | |
| # return prob | |