checkpointing

Author: lucidrains

alphafold3_pytorch/trainer.py

@@ -110,6 +110,8 @@ def __init__(
         default_lambda_lr = default_lambda_lr_fn,
         fabric: Fabric | None = None,
         accelerator = 'auto',
+        checkpoint_every: int = 1000,
+        checkpoint_folder: str = './checkpoints',
         fabric_kwargs: dict = dict(),
         ema_kwargs: dict = dict()
     ):
@@ -194,6 +196,14 @@ def __init__(
 
         self.steps = 0
 
+        # checkpointing logic
+
+        self.checkpoint_every = checkpoint_every
+        self.checkpoint_folder = Path(checkpoint_folder)
+
+        self.checkpoint_folder.mkdir(exist_ok = True, parents = True)
+        assert self.checkpoint_folder.is_dir()
+
     @property
     def is_main(self):
         return self.fabric.global_rank == 0
@@ -356,6 +366,11 @@ def __call__(
 
             self.wait()
 
+            if self.is_main and divisible_by(self.steps, self.checkpoint_every):
+                self.save(self.checkpoint_folder / f'af3.ckpt.{self.steps}.pt')
+
+            self.wait()
+
         # maybe test
 
         if self.is_main and self.needs_test:

pyproject.toml

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

tests/test_trainer.py

@@ -132,11 +132,14 @@ def test_trainer():
         num_train_steps = 2,
         batch_size = 1,
         valid_every = 1,
-        grad_accum_every = 2
+        grad_accum_every = 2,
+        checkpoint_every = 1
     )
 
     trainer()
 
+    assert Path('./checkpoints/af3.ckpt.1.pt').exists()
+
     # saving and loading from trainer
 
     trainer.save('./some/nested/folder2/training', overwrite = True)