fix tests

Author: lucidrains

.github/workflows/test.yml

@@ -6,6 +6,8 @@ jobs:
 
     runs-on: ubuntu-latest
     timeout-minutes: 20
+    strategy:
+      fail-fast: false
 
     steps:
     - uses: actions/checkout@v4

alphafold3_pytorch/trainer.py

@@ -267,6 +267,7 @@ def __init__(
         checkpoint_folder: str = './checkpoints',
         overwrite_checkpoints: bool = False,
         fabric_kwargs: dict = dict(),
+        use_ema: bool = True,
         ema_kwargs: dict = dict(
             use_foreach = True
         )
@@ -285,7 +286,10 @@ def __init__(
 
         # exponential moving average
 
-        if self.is_main:
+        self.ema_model = None
+        self.has_ema = self.is_main and use_ema
+
+        if self.has_ema:
             self.ema_model = EMA(
                 model,
                 beta = ema_decay,
@@ -574,7 +578,7 @@ def __call__(
 
             self.wait()
 
-            if self.is_main:
+            if self.has_ema:
                 self.ema_model.update()
 
             self.wait()
@@ -593,14 +597,16 @@ def __call__(
                 self.needs_valid and
                 divisible_by(self.steps, self.valid_every)
             ):
+                eval_model = default(self.ema_model, self.model)
+
                 with torch.no_grad():
-                    self.ema_model.eval()
+                    eval_model.eval()
 
                     total_valid_loss = 0.
                     valid_loss_breakdown = None
 
                     for valid_batch in self.valid_dataloader:
-                        valid_loss, loss_breakdown = self.ema_model(
+                        valid_loss, loss_breakdown = eval_model(
                             **valid_batch.dict(),
                             return_loss_breakdown = True
                         )
@@ -631,14 +637,16 @@ def __call__(
         # maybe test
 
         if self.is_main and self.needs_test:
+            eval_model = default(self.ema_model, self.model)
+
             with torch.no_grad():
-                self.ema_model.eval()
+                eval_model.eval()
 
                 total_test_loss = 0.
                 test_loss_breakdown = None
 
                 for test_batch in self.test_dataloader:
-                    test_loss, loss_breakdown = self.ema_model(
+                    test_loss, loss_breakdown = eval_model(
                         **test_batch.dict(),
                         return_loss_breakdown = True
                     )

tests/test_af3.py

@@ -504,7 +504,7 @@ def test_alphafold3(
     template_feats = torch.randn(2, 2, seq_len, seq_len, 44)
     template_mask = torch.ones((2, 2)).bool()
 
-    msa = torch.randn(2, 7, seq_len, 64)
+    msa = torch.randn(2, 7, seq_len, 8)
     msa_mask = torch.ones((2, 7)).bool()
 
     atom_pos = torch.randn(2, atom_seq_len, 3)
@@ -519,7 +519,9 @@ def test_alphafold3(
 
     alphafold3 = Alphafold3(
         dim_atom_inputs = 77,
-        dim_pairwise = 64,
+        dim_pairwise = 8,
+        dim_single = 8,
+        dim_token = 8,
         atoms_per_window = atoms_per_window,
         dim_template_feats = 44,
         num_dist_bins = 38,
@@ -531,15 +533,28 @@ def test_alphafold3(
             pairformer_stack_depth = 1
         ),
         msa_module_kwargs = dict(
-            depth = 1
+            depth = 1,
+            dim_msa = 8,
         ),
-        pairformer_stack = dict(
-            depth = 2
+        pairformer_stack=dict(
+            depth=1,
+            pair_bias_attn_dim_head = 4,
+            pair_bias_attn_heads = 2,
         ),
-        diffusion_module_kwargs = dict(
-            atom_encoder_depth = 1,
-            token_transformer_depth = 1,
-            atom_decoder_depth = 1,
+        diffusion_module_kwargs=dict(
+            atom_encoder_depth=1,
+            token_transformer_depth=1,
+            atom_decoder_depth=1,
+            atom_decoder_kwargs = dict(
+                attn_pair_bias_kwargs = dict(
+                    dim_head = 4
+                )
+            ),
+            atom_encoder_kwargs = dict(
+                attn_pair_bias_kwargs = dict(
+                    dim_head = 4
+                )
+            )
         ),
         stochastic_frame_average = stochastic_frame_average,
         confidence_head_atom_resolution = confidence_head_atom_resolution
@@ -569,6 +584,7 @@ def test_alphafold3(
         pde_labels = pde_labels,
         plddt_labels = plddt_labels,
         resolved_labels = resolved_labels,
+        num_rollout_steps = 1,
         diffusion_add_smooth_lddt_loss = True,
         return_loss_breakdown = True
     )

tests/test_input.py

@@ -174,11 +174,25 @@ def test_pdbinput_input():
         confidence_head_kwargs=dict(pairformer_depth=1),
         template_embedder_kwargs=dict(pairformer_stack_depth=1),
         msa_module_kwargs=dict(depth=1),
-        pairformer_stack=dict(depth=1),
+        pairformer_stack=dict(
+            depth=1,
+            pair_bias_attn_dim_head = 4,
+            pair_bias_attn_heads = 2,
+        ),
         diffusion_module_kwargs=dict(
             atom_encoder_depth=1,
             token_transformer_depth=1,
             atom_decoder_depth=1,
+            atom_decoder_kwargs = dict(
+                attn_pair_bias_kwargs = dict(
+                    dim_head = 4
+                )
+            ),
+            atom_encoder_kwargs = dict(
+                attn_pair_bias_kwargs = dict(
+                    dim_head = 4
+                )
+            )
         ),
     )
 

tests/test_trainer.py

@@ -223,7 +223,7 @@ def populate_mock_pdb_and_remove_test_folders():
     valid_folder.mkdir(exist_ok = True, parents = True)
     test_folder.mkdir(exist_ok = True, parents = True)
 
-    for i in range(100):
+    for i in range(2):
         shutil.copy2(str(working_cif_file), str(train_folder / f'{i}.cif'))
 
     for i in range(1):
@@ -287,7 +287,9 @@ def test_trainer_with_pdb_input(populate_mock_pdb_and_remove_test_folders):
     inputs = next(iter(dataloader))
 
     alphafold3.eval()
-    _, breakdown = alphafold3(**asdict(inputs), return_loss_breakdown = True)
+    with torch.no_grad():
+        _, breakdown = alphafold3(**asdict(inputs), return_loss_breakdown = True)
+
     before_distogram = breakdown.distogram
 
     path = './test-folder/nested/folder/af3'
@@ -298,7 +300,9 @@ def test_trainer_with_pdb_input(populate_mock_pdb_and_remove_test_folders):
     alphafold3 = Alphafold3.init_and_load(path)
 
     alphafold3.eval()
-    _, breakdown = alphafold3(**asdict(inputs), return_loss_breakdown = True)
+    with torch.no_grad():
+        _, breakdown = alphafold3(**asdict(inputs), return_loss_breakdown = True)
+
     after_distogram = breakdown.distogram
 
     assert torch.allclose(before_distogram, after_distogram)
@@ -318,6 +322,7 @@ def test_trainer_with_pdb_input(populate_mock_pdb_and_remove_test_folders):
         checkpoint_every = 1,
         checkpoint_folder = './test-folder/checkpoints',
         overwrite_checkpoints = True,
+        use_ema = False,
         ema_kwargs = dict(
             use_foreach = True,
             update_after_step = 0,