add test dataset

Author: lucidrains

alphafold3_pytorch/trainer.py

@@ -97,6 +97,7 @@ def __init__(
         grad_accum_every: int = 1,
         valid_dataset: Dataset | None = None,
         valid_every: int = 1000,
+        test_dataset: Dataset | None = None,
         optimizer: Optimizer | None = None,
         scheduler: LRScheduler | None = None,
         ema_decay = 0.999,
@@ -159,6 +160,14 @@ def __init__(
             self.valid_dataset_size = len(valid_dataset)
             self.valid_dataloader = DataLoader(valid_dataset, batch_size = batch_size)
 
+        # testing dataloader on EMA model
+
+        self.needs_test = exists(test_dataset)
+
+        if self.needs_test and self.is_main:
+            self.test_dataset_size = len(test_dataset)
+            self.test_dataloader = DataLoader(test_dataset, batch_size = batch_size)
+
         # training steps and num gradient accum steps
 
         self.num_train_steps = num_train_steps
@@ -347,4 +356,35 @@ def __call__(
 
             self.wait()
 
+        # maybe test
+
+        if self.is_main and self.needs_test:
+            with torch.no_grad():
+                self.ema_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_batch,
+                        return_loss_breakdown = True
+                    )
+
+                    test_batch_size = test_batch.get('atom_inputs').shape[0]
+                    scale = test_batch_size / self.test_dataset_size
+
+                    total_test_loss += test_loss.item() * scale
+                    test_loss_breakdown = accum_dict(test_loss_breakdown, loss_breakdown._asdict(), scale = scale)
+
+                self.print(f'test loss: {total_test_loss:.3f}')
+
+            # prepend test_ to all losses for logging
+
+            test_loss_breakdown = {f'test_{k}':v for k, v in test_loss_breakdown.items()}
+
+            # log
+
+            self.log(**test_loss_breakdown)
+
         print(f'training complete')

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "alphafold3-pytorch"
-version = "0.0.58"
+version = "0.0.59"
 description = "Alphafold 3 - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" }

tests/test_trainer.py

@@ -96,7 +96,8 @@ def test_trainer():
     )
 
     dataset = MockAtomDataset(100)
-    valid_dataset = MockAtomDataset(2)
+    valid_dataset = MockAtomDataset(4)
+    test_dataset = MockAtomDataset(2)
 
     # test saving and loading from Alphafold3, independent of lightning
 
@@ -126,6 +127,7 @@ def test_trainer():
         alphafold3,
         dataset = dataset,
         valid_dataset = valid_dataset,
+        test_dataset = test_dataset,
         accelerator = 'cpu',
         num_train_steps = 2,
         batch_size = 1,