Merge remote-tracking branch 'upstream/dev' into dev

Author: sfluegel

chebai/loggers/custom.py

@@ -1,5 +1,5 @@
 from datetime import datetime
-from typing import Literal, Optional, Union
+from typing import Literal, Optional, Union, List
 import os
 
 from lightning.fabric.utilities.types import _PATH
@@ -22,13 +22,16 @@ def __init__(
         entity: Optional[str] = None,
         offline: bool = False,
         log_model: Union[Literal["all"], bool] = False,
+        verbose_hyperparameters: bool = False,
+        tags: Optional[List[str]] = None,
         **kwargs,
     ):
         if version is None:
             version = f"{datetime.now():%y%m%d-%H%M}"
         self._version = version
         self._name = name
         self._fold = fold
+        self.verbose_hyperparameters = verbose_hyperparameters
         super().__init__(
             name=self.name,
             save_dir=save_dir,
@@ -40,6 +43,8 @@ def __init__(
             offline=offline,
             **kwargs,
         )
+        if tags:
+            self.experiment.tags += tuple(tags)
 
     @property
     def name(self) -> Optional[str]:

chebai/preprocessing/datasets/base.py

@@ -108,6 +108,7 @@ def __init__(
         if self.use_inner_cross_validation:
             os.makedirs(os.path.join(self.raw_dir, self.fold_dir), exist_ok=True)
             os.makedirs(os.path.join(self.processed_dir, self.fold_dir), exist_ok=True)
+        self.save_hyperparameters()
 
     @property
     def identifier(self):

chebai/trainer/CustomTrainer.py

@@ -9,6 +9,7 @@
 import torch
 
 from chebai.preprocessing.reader import CLS_TOKEN, ChemDataReader
+from chebai.loggers.custom import CustomLogger
 
 log = logging.getLogger(__name__)
 
@@ -20,6 +21,35 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         # instantiation custom logger connector
         self._logger_connector.on_trainer_init(self.logger, 1)
+        # log additional hyperparameters to wandb
+        if isinstance(self.logger, CustomLogger):
+            custom_logger = self.logger
+            assert isinstance(custom_logger, CustomLogger)
+            if custom_logger.verbose_hyperparameters:
+                log_kwargs = {}
+                for key, value in self.init_kwargs.items():
+                    log_key, log_value = self._resolve_logging_argument(key, value)
+                    log_kwargs[log_key] = log_value
+                self.logger.log_hyperparams(log_kwargs)
+
+    def _resolve_logging_argument(self, key, value):
+        if isinstance(value, list):
+            key_value_pairs = [
+                self._resolve_logging_argument(f"{key}_{i}", v)
+                for i, v in enumerate(value)
+            ]
+            return key, {k: v for k, v in key_value_pairs}
+        if not (
+            isinstance(value, str)
+            or isinstance(value, float)
+            or isinstance(value, int)
+            or value is None
+        ):
+            params = {"class": value.__class__}
+            params.update(value.__dict__)
+            return key, params
+        else:
+            return key, value
 
     def predict_from_file(
         self,