ready for fine-tuning

Author: lucidrains

alphafold3_pytorch/trainer.py

@@ -354,7 +354,8 @@ def load(
         self,
         path: str | Path,
         strict = True,
-        prefix = None
+        prefix = None,
+        only_model = False
     ):
         if isinstance(path, str):
             path = Path(path)
@@ -382,6 +383,11 @@ def load(
 
         self.model.load(path)
 
+        if only_model:
+            return
+
+        # load optimizer and scheduler states
+
         package = torch.load(str(path))
 
         if 'optimizer' in package:

pyproject.toml

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

tests/test_trainer.py

@@ -164,6 +164,10 @@ def test_trainer():
     trainer.save('./some/nested/folder2/training.pt', overwrite = True)
     trainer.load('./some/nested/folder2/training.pt')
 
+    # allow for only loading model, needed for fine-tuning logic
+
+    trainer.load('./some/nested/folder2/training.pt', only_model = True)
+
     # also allow for loading Alphafold3 directly from training ckpt
 
     alphafold3 = Alphafold3.init_and_load('./some/nested/folder2/training.pt')