add ability to instantiate trainer directly from yaml file, only needing to pass in datasets

Author: lucidrains

alphafold3_pytorch/__init__.py

@@ -39,7 +39,8 @@
 )
 
 from alphafold3_pytorch.configs import (
-    Alphafold3Config
+    Alphafold3Config,
+    TrainerConfig
 )
 
 __all__ = [
@@ -72,5 +73,6 @@
     Alphafold3,
     Alphafold3Config,
     AtomInput,
-    Trainer
+    Trainer,
+    TrainerConfig
 ]

alphafold3_pytorch/alphafold3.py

@@ -17,7 +17,7 @@
     Sequential,
 )
 
-from typing import Literal, Tuple, NamedTuple, Callable
+from typing import List, Literal, Tuple, NamedTuple, Callable
 
 from alphafold3_pytorch.typing import (
     Float,
@@ -2830,7 +2830,7 @@ def __init__(
         dim_single = 384,
         dim_pairwise = 128,
         dim_token = 768,
-        distance_bins: List[float] = torch.linspace(3, 20, 38).tolist(),
+        distance_bins: List[float] = torch.linspace(3, 20, 38).float().tolist(),
         ignore_index = -1,
         num_dist_bins: int | None = None,
         num_plddt_bins = 50,
@@ -3020,12 +3020,12 @@ def __init__(
 
         # logit heads
 
-        distance_bins = Tensor(distance_bins)
+        distance_bins_tensor = Tensor(distance_bins)
 
-        self.register_buffer('distance_bins', distance_bins)
-        num_dist_bins = default(num_dist_bins, len(distance_bins))
+        self.register_buffer('distance_bins', distance_bins_tensor)
+        num_dist_bins = default(num_dist_bins, len(distance_bins_tensor))
 
-        assert len(distance_bins) == num_dist_bins, '`distance_bins` must have a length equal to the `num_dist_bins` passed in'
+        assert len(distance_bins_tensor) == num_dist_bins, '`distance_bins` must have a length equal to the `num_dist_bins` passed in'
 
         self.distogram_head = DistogramHead(
             dim_pairwise = dim_pairwise,

alphafold3_pytorch/configs.py

@@ -1,8 +1,18 @@
 from __future__ import annotations
 
 from alphafold3_pytorch.typing import typecheck
+from typing import Callable, List
+
 from alphafold3_pytorch.alphafold3 import Alphafold3
 
+from alphafold3_pytorch.trainer import (
+    Trainer,
+    Dataset,
+    Fabric,
+    Optimizer,
+    LRScheduler
+)
+
 import yaml
 from pathlib import Path
 
@@ -27,7 +37,7 @@ def yaml_config_path_to_dict(
         maybe_config_dict = yaml.safe_load(f)
 
     assert exists(maybe_config_dict), f'unable to parse yaml config at {str(path)}'
-    assert isinstance(maybe_config_dict, dict), f'yaml config file is not a dictionary'
+    assert isinstance(maybe_config_dict, dict), 'yaml config file is not a dictionary'
 
     return maybe_config_dict
 
@@ -76,4 +86,59 @@ def create_instance_from_yaml_file(path: str | Path) -> Alphafold3:
         return af3_config.create_instance()
 
 class TrainerConfig(BaseModelWithExtra):
-    pass
+    model: Alphafold3Config
+    num_train_steps: int
+    batch_size: int
+    grad_accum_every: int
+    valid_every: int
+    ema_decay: float
+    lr: float
+    clip_grad_norm: int | float
+    accelerator: str 
+    checkpoint_prefix: str
+    checkpoint_every: int
+    checkpoint_folder: str
+    overwrite_checkpoints: bool
+
+    @staticmethod
+    @typecheck
+    def from_yaml_file(path: str | Path):
+        config_dict = yaml_config_path_to_dict(path)
+        return TrainerConfig(**config_dict)
+
+    def create_instance(
+        self,
+        dataset: Dataset,
+        fabric: Fabric | None = None,
+        test_dataset: Dataset | None = None,
+        optimizer: Optimizer | None = None,
+        scheduler: LRScheduler | None = None,
+        valid_dataset: Dataset | None = None,
+        map_dataset_input_fn: Callable | None = None,
+    ) -> Trainer:
+
+        trainer_kwargs = self.model_dump()
+
+        alphafold3 = self.model.create_instance()
+
+        trainer_kwargs.update(dict(
+            model = alphafold3,
+            dataset = dataset,
+            fabric = fabric,
+            test_dataset = test_dataset,
+            optimizer = optimizer,
+            scheduler = scheduler,
+            valid_dataset = valid_dataset,
+            map_dataset_input_fn = map_dataset_input_fn
+        ))
+
+        trainer = Trainer(**trainer_kwargs)
+        return trainer
+
+    def create_instance_from_yaml_file(
+        path: str | Path,
+        **kwargs
+    ) -> Trainer:
+
+        trainer_config = TrainerConfig.from_yaml_file(path)
+        return trainer_config.create_instance(**kwargs)

alphafold3_pytorch/pdb_dataset_curation.py

@@ -188,7 +188,7 @@
 
 # Helper functions
 
-def exists(v: Any) -> bool:
+def exists(v) -> bool:
     """Return `False` if `v` is `None`, else return `True`."""
     return v is not None
 

pyproject.toml

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

tests/alphafold3.yaml

@@ -10,7 +10,6 @@ dim_input_embedder_token: 384
 dim_single: 384
 dim_pairwise: 128
 dim_token: 768
-distance_bins: 38
 ignore_index: -1
 num_dist_bins: null
 num_plddt_bins: 50

tests/test_af3.py

@@ -3,6 +3,7 @@
 
 import torch
 import pytest
+from pathlib import Path
 
 from alphafold3_pytorch import (
     SmoothLDDTLoss,
@@ -22,6 +23,7 @@
     ConfidenceHead,
     DistogramHead,
     Alphafold3,
+    Alphafold3Config
 )
 
 from alphafold3_pytorch.alphafold3 import (
@@ -361,7 +363,7 @@ def test_confidence_head():
 
     confidence_head = ConfidenceHead(
         dim_single_inputs = 77,
-        atompair_dist_bins = torch.linspace(3, 20, 37),
+        atompair_dist_bins = torch.linspace(3, 20, 37).tolist(),
         dim_single = 384,
         dim_pairwise = 128,
     )
@@ -565,3 +567,12 @@ def test_alphafold3_without_msa_and_templates():
     )
 
     loss.backward()
+
+# test creation from config
+
+def test_alphafold3_config():
+    curr_dir = Path(__file__).parents[0]
+    af3_yaml = curr_dir / 'alphafold3.yaml'
+
+    alphafold3 = Alphafold3Config.create_instance_from_yaml_file(af3_yaml)
+    assert isinstance(alphafold3, Alphafold3)

tests/test_config.py

@@ -12,7 +12,8 @@
     Alphafold3,
     AtomInput,
     DataLoader,
-    Trainer
+    Trainer,
+    TrainerConfig
 )
 
 # mock dataset
@@ -165,3 +166,20 @@ def test_trainer():
     # also allow for loading Alphafold3 directly from training ckpt
 
     alphafold3 = Alphafold3.init_and_load('./some/nested/folder2/training.pt')
+
+# test creating trainer + alphafold3 from config
+
+def test_trainer_config():
+    curr_dir = Path(__file__).parents[0]
+    trainer_yaml_path = curr_dir / 'trainer.yaml'
+
+    trainer = TrainerConfig.create_instance_from_yaml_file(
+        trainer_yaml_path,
+        dataset = MockAtomDataset(16)
+    )
+
+    assert isinstance(trainer, Trainer)
+
+    # take a single training step
+
+    trainer()

tests/test_trainer.py

@@ -0,0 +1,35 @@
+---
+model:
+    dim_atom_inputs: 77
+    dim_template_feats: 44
+    dim_template_model: 64
+    atoms_per_window: 27
+    dim_atom: 128
+    dim_atompair_inputs: 5
+    dim_atompair: 16
+    dim_input_embedder_token: 384
+    dim_single: 384
+    dim_pairwise: 128
+    dim_token: 768
+    ignore_index: -1
+    num_dist_bins: null
+    num_plddt_bins: 50
+    num_pde_bins: 64
+    num_pae_bins: 64
+    sigma_data: 16
+    diffusion_num_augmentations: 4
+    loss_confidence_weight: 0.0001
+    loss_distogram_weight: 0.01
+    loss_diffusion_weight: 4.
+num_train_steps: 1
+batch_size: 1
+grad_accum_every: 1
+valid_every: 1
+ema_decay: 0.999
+lr: 0.0001
+clip_grad_norm: 10.
+accelerator: cpu
+checkpoint_prefix: af3.ckpt.
+checkpoint_every: 1000
+checkpoint_folder: ./checkpoints
+overwrite_checkpoints: false