handle a directory being passed into Trainer.load

Author: lucidrains

alphafold3_pytorch/trainer.py

@@ -316,6 +316,10 @@ def __init__(
         self.checkpoint_folder.mkdir(exist_ok = True, parents = True)
         assert self.checkpoint_folder.is_dir()
 
+        # save the path for the last loaded model, if any
+
+        self.model_loaded_from_path = None
+
     @property
     def is_main(self):
         return self.fabric.global_rank == 0
@@ -345,6 +349,23 @@ def load(self, path: str | Path, strict = True):
 
         assert path.exists()
 
+        # if the path is a directory, then automatically load latest checkpoint
+
+        if path.is_dir():
+            model_paths = [*path.glob('**/*.pt')]
+
+            assert len(model_paths) > 0, f'no files found in directory {path}'
+
+            model_paths = sorted(model_paths, key = lambda p: int(str(p).split('.')[-2]))
+
+            path = model_paths[-1]
+
+        # for eventually saving entire training history in filename
+
+        self.model_loaded_from_path = path
+
+        # load model from path
+
         self.model.load(path)
 
         package = torch.load(str(path))

pyproject.toml

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

tests/test_trainer.py

@@ -94,7 +94,7 @@ def test_trainer():
             depth = 1
         ),
         pairformer_stack = dict(
-            depth = 2
+            depth = 1
         ),
         diffusion_module_kwargs = dict(
             atom_encoder_depth = 1,
@@ -147,13 +147,21 @@ def test_trainer():
 
     trainer()
 
+    # assert checkpoints created
+
     assert Path('./checkpoints/af3.ckpt.1.pt').exists()
 
+    # assert can load latest checkpoint by loading from a directory
+
+    trainer.load('./checkpoints')
+
+    assert str(trainer.model_loaded_from_path) == str(Path('./checkpoints/af3.ckpt.2.pt'))
+
     # saving and loading from trainer
 
-    trainer.save('./some/nested/folder2/training', overwrite = True)
-    trainer.load('./some/nested/folder2/training')
+    trainer.save('./some/nested/folder2/training.pt', overwrite = True)
+    trainer.load('./some/nested/folder2/training.pt')
 
     # also allow for loading Alphafold3 directly from training ckpt
 
-    alphafold3 = Alphafold3.init_and_load('./some/nested/folder2/training')
+    alphafold3 = Alphafold3.init_and_load('./some/nested/folder2/training.pt')