(self, d_msa=256, d_pair=128, n_head=8, d_state=16,
d_hidden=32, p_drop=0.15, use_global_attn=False)
| 12 | # Update MSA with biased self-attention. bias from Pair & Str |
| 13 | class MSAPairStr2MSA(nn.Module): |
| 14 | def __init__(self, d_msa=256, d_pair=128, n_head=8, d_state=16, |
| 15 | d_hidden=32, p_drop=0.15, use_global_attn=False): |
| 16 | super(MSAPairStr2MSA, self).__init__() |
| 17 | self.norm_pair = nn.LayerNorm(d_pair) |
| 18 | self.proj_pair = nn.Linear(d_pair+36, d_pair) |
| 19 | self.norm_state = nn.LayerNorm(d_state) |
| 20 | self.proj_state = nn.Linear(d_state, d_msa) |
| 21 | self.drop_row = Dropout(broadcast_dim=1, p_drop=p_drop) |
| 22 | self.row_attn = MSARowAttentionWithBias(d_msa=d_msa, d_pair=d_pair, |
| 23 | n_head=n_head, d_hidden=d_hidden) |
| 24 | if use_global_attn: |
| 25 | self.col_attn = MSAColGlobalAttention(d_msa=d_msa, n_head=n_head, d_hidden=d_hidden) |
| 26 | else: |
| 27 | self.col_attn = MSAColAttention(d_msa=d_msa, n_head=n_head, d_hidden=d_hidden) |
| 28 | self.ff = FeedForwardLayer(d_msa, 4, p_drop=p_drop) |
| 29 | |
| 30 | # Do proper initialization |
| 31 | self.reset_parameter() |
| 32 | |
| 33 | def reset_parameter(self): |
| 34 | # initialize weights to normal distrib |
nothing calls this directly
no test coverage detected