| 206 | return self.drop(msa) |
| 207 | |
| 208 | class TemplatePairStack(nn.Module): |
| 209 | # process template pairwise features |
| 210 | # use structure-biased attention |
| 211 | def __init__(self, n_block=2, d_templ=64, n_head=4, d_hidden=16, p_drop=0.25): |
| 212 | super(TemplatePairStack, self).__init__() |
| 213 | self.n_block = n_block |
| 214 | proc_s = [PairStr2Pair(d_pair=d_templ, n_head=n_head, d_hidden=d_hidden, p_drop=p_drop) for i in range(n_block)] |
| 215 | self.block = nn.ModuleList(proc_s) |
| 216 | self.norm = nn.LayerNorm(d_templ) |
| 217 | def forward(self, templ, rbf_feat, use_checkpoint=False): |
| 218 | B, T, L = templ.shape[:3] |
| 219 | templ = templ.reshape(B*T, L, L, -1) |
| 220 | |
| 221 | for i_block in range(self.n_block): |
| 222 | if use_checkpoint: |
| 223 | templ = checkpoint.checkpoint(create_custom_forward(self.block[i_block]), templ, rbf_feat) |
| 224 | else: |
| 225 | templ = self.block[i_block](templ, rbf_feat) |
| 226 | return self.norm(templ).reshape(B, T, L, L, -1) |
| 227 | |
| 228 | class TemplateTorsionStack(nn.Module): |
| 229 | def __init__(self, n_block=2, d_templ=64, n_head=4, d_hidden=16, p_drop=0.15): |