(self, n_extra_block=4, n_main_block=12, n_ref_block=4,
d_msa=256, d_msa_full=64, d_pair=128, d_hidden=32,
n_head_msa=8, n_head_pair=4,
SE3_param_full={'l0_in_features':32, 'l0_out_features':16, 'num_edge_features':32},
SE3_param_topk={'l0_in_features':32, 'l0_out_features':16, 'num_edge_features':32},
p_drop=0.15)
| 335 | |
| 336 | class IterativeSimulator(nn.Module): |
| 337 | def __init__(self, n_extra_block=4, n_main_block=12, n_ref_block=4, |
| 338 | d_msa=256, d_msa_full=64, d_pair=128, d_hidden=32, |
| 339 | n_head_msa=8, n_head_pair=4, |
| 340 | SE3_param_full={'l0_in_features':32, 'l0_out_features':16, 'num_edge_features':32}, |
| 341 | SE3_param_topk={'l0_in_features':32, 'l0_out_features':16, 'num_edge_features':32}, |
| 342 | p_drop=0.15): |
| 343 | super(IterativeSimulator, self).__init__() |
| 344 | self.n_extra_block = n_extra_block |
| 345 | self.n_main_block = n_main_block |
| 346 | self.n_ref_block = n_ref_block |
| 347 | |
| 348 | self.proj_state = nn.Linear(SE3_param_topk['l0_out_features'], SE3_param_full['l0_out_features']) |
| 349 | # Update with extra sequences |
| 350 | if n_extra_block > 0: |
| 351 | self.extra_block = nn.ModuleList([IterBlock(d_msa=d_msa_full, d_pair=d_pair, |
| 352 | n_head_msa=n_head_msa, |
| 353 | n_head_pair=n_head_pair, |
| 354 | d_hidden_msa=8, |
| 355 | d_hidden=d_hidden, |
| 356 | p_drop=p_drop, |
| 357 | use_global_attn=True, |
| 358 | SE3_param=SE3_param_full) |
| 359 | for i in range(n_extra_block)]) |
| 360 | |
| 361 | # Update with seed sequences |
| 362 | if n_main_block > 0: |
| 363 | self.main_block = nn.ModuleList([IterBlock(d_msa=d_msa, d_pair=d_pair, |
| 364 | n_head_msa=n_head_msa, |
| 365 | n_head_pair=n_head_pair, |
| 366 | d_hidden=d_hidden, |
| 367 | p_drop=p_drop, |
| 368 | use_global_attn=False, |
| 369 | SE3_param=SE3_param_full) |
| 370 | for i in range(n_main_block)]) |
| 371 | |
| 372 | self.proj_state2 = nn.Linear(SE3_param_full['l0_out_features'], SE3_param_topk['l0_out_features']) |
| 373 | # Final SE(3) refinement |
| 374 | if n_ref_block > 0: |
| 375 | self.str_refiner = Str2Str(d_msa=d_msa, d_pair=d_pair, |
| 376 | d_state=SE3_param_topk['l0_out_features'], |
| 377 | SE3_param=SE3_param_topk, |
| 378 | p_drop=p_drop) |
| 379 | |
| 380 | self.reset_parameter() |
| 381 | def reset_parameter(self): |
| 382 | self.proj_state = init_lecun_normal(self.proj_state) |
| 383 | nn.init.zeros_(self.proj_state.bias) |
nothing calls this directly
no test coverage detected