default to tensorboard logging, add wandb later

Author: lucidrains

.gitignore

@@ -1,4 +1,5 @@
 .env
+logs/
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

alphafold3_pytorch/configs.py

@@ -31,6 +31,8 @@
     DirectoryPath
 )
 
+from lightning.fabric.loggers import TensorBoardLogger
+
 # functions
 
 def exists(v):
@@ -166,6 +168,9 @@ class TrainerConfig(BaseModelWithExtra):
     checkpoint_folder: str
     overwrite_checkpoints: bool
     dataset_config: DatasetConfig | None = None
+    use_tensorboard: bool = True
+    tensorboard_log_dir: str = './logs'
+    logger_kwargs: dict = dict()
 
     @classmethod
     @typecheck
@@ -193,7 +198,12 @@ def create_instance(
     ) -> Trainer:
 
         trainer_kwargs = self.model_dump(
-            exclude = {'dataset_config'}
+            exclude = {
+                'dataset_config',
+                'use_tensorboard',
+                'tensorboard_log_dir',
+                'logger_kwargs'
+            }
         )
 
         assert exists(self.model) ^ exists(model), 'either model is available on the trainer config, or passed in when creating the instance, but not both or neither'
@@ -261,6 +271,13 @@ def create_instance(
 
         assert 'dataset' in trainer_kwargs, 'dataset is absent - dataset_type must be specified along with train folders (pdb for now), or the Dataset instance must be passed in'
 
+        # handle loggers
+
+        loggers = []
+
+        if self.use_tensorboard:
+            loggers.append(TensorBoardLogger(self.tensorboard_log_dir, **self.logger_kwargs))
+
         # handle rest
 
         trainer_kwargs.update(dict(
@@ -270,7 +287,8 @@ def create_instance(
             optimizer = optimizer,
             scheduler = scheduler,
             valid_dataset = valid_dataset,
-            map_dataset_input_fn = map_dataset_input_fn
+            map_dataset_input_fn = map_dataset_input_fn,
+            loggers = loggers
         ))
 
         trainer = Trainer(**trainer_kwargs)

alphafold3_pytorch/trainer.py

@@ -47,6 +47,7 @@
 from ema_pytorch import EMA
 
 from lightning import Fabric
+from lightning.fabric.loggers import Logger
 from lightning.fabric.wrappers import _unwrap_objects
 
 from shortuuid import uuid
@@ -290,6 +291,7 @@ def __init__(
         default_lambda_lr = default_lambda_lr_fn,
         train_sampler: Sampler | None = None,
         fabric: Fabric | None = None,
+        loggers: List[Logger] = [],
         accelerator = 'auto',
         checkpoint_prefix = 'af3.ckpt.',
         checkpoint_every: int = 1000,
@@ -317,7 +319,11 @@ def __init__(
         # instantiate fabric
 
         if not exists(fabric):
-            fabric = Fabric(accelerator = accelerator, **fabric_kwargs)
+            fabric = Fabric(
+                accelerator = accelerator,
+                loggers = loggers,
+                **fabric_kwargs
+            )
 
         self.fabric = fabric
         fabric.launch()

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "alphafold3-pytorch"
-version = "0.2.111"
+version = "0.2.114"
 description = "Alphafold 3 - Pytorch"
 authors = [
     { name = "Phil Wang", email = "[email protected]" },
@@ -48,6 +48,7 @@ dependencies = [
     "scikit-learn>=1.5.0",
     "sh>=2.0.7",
     "shortuuid",
+    "tensorboard",
     "taylor-series-linear-attention>=0.1.9",
     "timeout_decorator>=0.5.0",
     'torch_geometric',