forked from Ligo-Biosciences/AlphaFold3
-
Notifications
You must be signed in to change notification settings - Fork 0
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
ardagoreci
committed
Aug 10, 2024
1 parent
12887c4
commit 732ebe0
Showing
21 changed files
with
793 additions
and
304 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,173 @@ | ||
| # AlphaFold3 configs | ||
| # _target_: src.models.model_wrapper.AlphaFoldWrapper | ||
| # config: | ||
| optimizer: | ||
| _target_: torch.optim.Adam # deepspeed.ops.adam.FusedAdam | ||
| _partial_: true | ||
| lr: 0.0018 # 0.0018 | ||
| betas: | ||
| - 0.9 | ||
| - 0.95 | ||
| eps: 1e-08 | ||
| weight_decay: 0.0 | ||
|
|
||
| scheduler: | ||
| _target_: src.utils.lr_schedulers.AlphaFoldLRScheduler | ||
| _partial_: true | ||
| last_epoch: -1 | ||
| verbose: false | ||
| base_lr: 0.0 # the starting learning rate | ||
| max_lr: 0.0018 | ||
| warmup_no_steps: 1000 | ||
| start_decay_after_n_steps: 50_000 | ||
| decay_every_n_steps: 50_000 | ||
| decay_factor: 0.95 | ||
|
|
||
| # Loss configs | ||
| loss: | ||
| diffusion_loss: | ||
| # mse: | ||
| # weight_dna: 5.0 | ||
| # weight_rna: 5.0 | ||
| # weight_ligand: 10.0 | ||
| # weight_protein: 1.0 | ||
| # smooth_lddt: | ||
| # weight: 1.0 | ||
| use_smooth_lddt: false | ||
| sd_data: 16.0 | ||
| weight: 1.0 | ||
|
|
||
| distogram: | ||
| min_bin: 2.3125 | ||
| max_bin: 21.6875 | ||
| no_bins: 64 | ||
| eps: 0.000006 # 1e-6 | ||
| weight: 0.03 | ||
|
|
||
| experimentally_resolved: | ||
| eps: 0.00000001 # 1e-8, | ||
| # min_resolution: 0.1, | ||
| # max_resolution: 3.0, | ||
| weight: 0.0004 | ||
|
|
||
| plddt_loss: | ||
| min_resolution: 0.1 | ||
| max_resolution: 3.0 | ||
| cutoff: 15.0 | ||
| no_bins": 50 | ||
| eps: 0.0000000001 # 1e-10, | ||
| weight: 0.0004 | ||
|
|
||
|
|
||
| # TODO: fix the model.model notation in interpolation | ||
| model: | ||
| c_token: 64 # the token representation dim | ||
| c_pair: 16 # the pair representation dim | ||
| c_atom: 16 # the atom representation dim | ||
| c_atompair: 16 # the atom pair representation dim | ||
|
|
||
| # Pair stack parameters (used in Pairformer, MSA module, and Confidence head) | ||
| c_hidden_tri_mul: 16 # the hidden dim for the triangle multiplicative update | ||
| c_hidden_pair_attn: 16 # the hidden dim for the pair attention ${common.c_hidden_pair_attn} | ||
| no_heads_tri_attn: 1 | ||
| transition_n: 1 | ||
| pair_dropout: 0.25 | ||
| fuse_projection_weights: false | ||
| blocks_per_ckpt: null # number of blocks per checkpoint, if none, no checkpointing | ||
| clear_cache_between_blocks: false # whether to clear GPU memory cache between blocks | ||
| # Pairformer attention pair bias | ||
| no_heads_single_attn: 1 | ||
|
|
||
| # Input Embedder | ||
| input_embedder: | ||
| c_token: ${model.model.c_token} | ||
| c_trunk_pair: ${model.model.c_pair} | ||
| c_atom: ${model.model.c_atom} | ||
| c_atompair: ${model.model.c_atompair} | ||
|
|
||
| # MSA module | ||
| msa_module: | ||
| no_blocks: 1 # 4 | ||
| c_msa: 16 | ||
| c_token: ${model.model.c_token} | ||
| c_z: ${model.model.c_pair} | ||
| c_hidden: 8 | ||
| no_heads: 1 | ||
| c_hidden_tri_mul: ${model.model.c_hidden_tri_mul} | ||
| c_hidden_pair_attn: ${model.model.c_hidden_pair_attn} | ||
| no_heads_tri_attn: ${model.model.no_heads_tri_attn} | ||
| transition_n: ${model.model.transition_n} | ||
| pair_dropout: ${model.model.pair_dropout} | ||
| fuse_projection_weights: ${model.model.fuse_projection_weights} | ||
| clear_cache_between_blocks: ${model.model.clear_cache_between_blocks} | ||
| blocks_per_ckpt: ${model.model.blocks_per_ckpt} | ||
| inf: 1e8 | ||
|
|
||
| # PairformerStack | ||
| pairformer_stack: | ||
| c_s: ${model.model.c_token} | ||
| c_z: ${model.model.c_pair} | ||
| no_blocks: 1 # 48 | ||
| c_hidden_mul: ${model.model.c_hidden_tri_mul} | ||
| c_hidden_pair_attn: ${model.model.c_hidden_pair_attn} | ||
| no_heads_tri_attn: ${model.model.no_heads_tri_attn} | ||
| no_heads_single_attn: ${model.model.no_heads_single_attn} | ||
| transition_n: ${model.model.transition_n} | ||
| pair_dropout: ${model.model.pair_dropout} | ||
| fuse_projection_weights: ${model.model.fuse_projection_weights} | ||
| blocks_per_ckpt: ${model.model.blocks_per_ckpt} | ||
| clear_cache_between_blocks: false | ||
| inf: 1e8 | ||
|
|
||
| # Diffusion module | ||
| diffusion_module: | ||
| c_atom: ${model.model.c_atom} | ||
| c_atompair: ${model.model.c_atompair} | ||
| c_token: ${model.model.c_token} | ||
| c_tokenpair: ${model.model.c_pair} | ||
| atom_encoder_blocks: 1 | ||
| atom_encoder_heads: 1 | ||
| dropout: 0.0 | ||
| atom_attention_n_queries: 32 # TODO: with sliding window attention this is not used. | ||
| atom_attention_n_keys: 128 | ||
| atom_decoder_blocks: 1 | ||
| atom_decoder_heads: 1 | ||
| token_transformer_blocks: 1 # 24 | ||
| token_transformer_heads: 1 | ||
| sd_data: 16.0 | ||
| s_max: 160.0 | ||
| s_min: 4e-4 | ||
| p: 7.0 | ||
| clear_cache_between_blocks: ${model.model.clear_cache_between_blocks} | ||
| blocks_per_ckpt: ${model.model.blocks_per_ckpt} | ||
| compile_module: false # TODO: this parameter does not belong here, will be removed in future iterations. | ||
|
|
||
| confidence_head: | ||
| c_s: 384 # ${model.c_token} | ||
| c_z: ${model.model.c_pair} | ||
| no_blocks: 4 | ||
| no_bins_pde: 64 | ||
| no_bins_plddt: 64 | ||
| no_bins_pae: 64 | ||
| c_hidden_mul: ${model.model.c_hidden_tri_mul} | ||
| c_hidden_pair_attn: ${model.model.c_hidden_pair_attn} | ||
| no_heads_tri_attn: ${model.model.no_heads_tri_attn} | ||
| no_heads_single_attn: ${model.model.no_heads_single_attn} | ||
| transition_n: ${model.model.transition_n} | ||
| pair_dropout: ${model.model.pair_dropout} | ||
| fuse_projection_weights: ${model.model.fuse_projection_weights} | ||
|
|
||
| distogram_head: | ||
| c_z: ${model.model.c_pair} | ||
| no_bins: 64 | ||
|
|
||
| # Exponential moving average decay rate | ||
| ema_decay: 0.999 | ||
|
|
||
| globals: | ||
| chunk_size: null # 4 | ||
| # Use DeepSpeed memory-efficient attention kernel in supported modules. | ||
| use_deepspeed_evo_attention: true | ||
| samples_per_trunk: 48 # Number of diffusion module replicas per trunk | ||
| rollout_samples_per_trunk: 1 # Number of mini rollouts per trunk | ||
| eps: 0.00000001 |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.