disable templates

Ligo-Biosciences · Aug 13, 2024 · 41ad3df · 41ad3df
1 parent bbdb23d
commit 41ad3df
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 18 deletions.
diff --git a/configs/data/erebor.yaml b/configs/data/erebor.yaml
@@ -100,7 +100,7 @@ config:
   data_module:
     use_small_bfd: false
     data_loaders:
-        batch_size: 1
+        batch_size: 2
         num_workers: 16
         pin_memory: true
 
@@ -220,7 +220,7 @@ config:
       - "between_segment_residues"
       - "deletion_matrix"
       - "no_recycling_iters"
-    use_templates: true
+    use_templates: false
     use_template_torsion_angles: false
     template_features:
       - "template_all_atom_positions"

diff --git a/configs/model/alphafold3.yaml b/configs/model/alphafold3.yaml
@@ -114,7 +114,7 @@ model:
   pairformer_stack:
     c_s: ${model.model.c_token}
     c_z: ${model.model.c_pair}
-    no_blocks: 48
+    no_blocks: 24  # 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}
@@ -139,7 +139,7 @@ model:
     atom_attention_n_keys: 128
     atom_decoder_blocks: 3
     atom_decoder_heads: 16
-    token_transformer_blocks: 24 # 24
+    token_transformer_blocks: 12 # 24
     token_transformer_heads: 16
     sd_data: 16.0
     s_max: 160.0

diff --git a/src/models/model.py b/src/models/model.py
@@ -22,9 +22,9 @@ def __init__(self, config):
             **self.config["input_embedder"]
         )
 
-        self.template_embedder = TemplateEmbedder(
-            **self.config["template_embedder"]
-        )
+        # self.template_embedder = TemplateEmbedder(
+        #    **self.config["template_embedder"]
+        # )
 
         self.msa_module = MSAModule(
             **self.config["msa_module"]
@@ -107,17 +107,17 @@ def run_trunk(
         del s_prev, z_prev, s_init, z_init
 
         # Embed the templates
-        z = add(z,
-                self.template_embedder(
-                    feats,
-                    z,
-                    pair_mask=pair_mask,
-                    chunk_size=self.globals.chunk_size,
-                    use_deepspeed_evo_attention=self.globals.use_deepspeed_evo_attention,
-                    inplace_safe=inplace_safe
-                ),
-                inplace=inplace_safe
-                )
+        # z = add(z,
+        #        self.template_embedder(
+        #            feats,
+        #            z,
+        #            pair_mask=pair_mask,
+        #            chunk_size=self.globals.chunk_size,
+        #            use_deepspeed_evo_attention=self.globals.use_deepspeed_evo_attention,
+        #            inplace_safe=inplace_safe
+        #        ),
+        #        inplace=inplace_safe
+        #        )
 
         # Process the MSA
         z = add(z,