add hyperparameter saving to datamodule and trainer

Author: sfluegel

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,7 +21,35 @@ def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         # instantiation custom logger connector
         self._logger_connector.on_trainer_init(self.logger, 1)
-        self.logger.log_hyperparams(self.init_kwargs)
+        # 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,