Merging ModelTrainer and TrainingAdapter class (#225)

* add tags to wandb logger
* replace experimental_name placehodler
* combine Trainer and TrainerAdaptor class
* update tests, names and class logic
* update docstrings, removed unsused methods
* change `model_name` to `potential_name`
* Fixed minor naming error; replaced Fire automatic command line generation with manually set up command line interface using argparse.
* changing model_seed to potential_seed; added docstrings
* setting random number seeds
* If neither runtime paramters and simulation environment, default to pytorch
* Added potential_seed to potential parameters

---------

Co-authored-by: chrisiacovella <[email protected]>

modelforge/potential/ani.py

@@ -650,6 +650,8 @@ class ANI2x(BaseNetwork):
         Configuration for postprocessing parameters.
     dataset_statistic : Optional[Dict[str, float]], optional
         Statistics of the dataset, by default None.
+    potential_seed : Optional[int], optional
+        Seed for the random number generator, default None.
     """
 
     def __init__(
@@ -664,6 +666,7 @@ def __init__(
         activation_function_parameter: Dict,
         postprocessing_parameter: Dict[str, Dict[str, bool]],
         dataset_statistic: Optional[Dict[str, float]] = None,
+        potential_seed: Optional[int] = None,
     ) -> None:
 
         from modelforge.utils.units import _convert_str_to_unit
@@ -674,6 +677,7 @@ def __init__(
             dataset_statistic=dataset_statistic,
             postprocessing_parameter=postprocessing_parameter,
             maximum_interaction_radius=_convert_str_to_unit(maximum_interaction_radius),
+            potential_seed=potential_seed,
         )
 
         activation_function = activation_function_parameter["activation_function"]

modelforge/potential/models.py

@@ -540,19 +540,41 @@ def apply_bwd(res, grads):
         return apply, model_params, model_buffer
 
 
+from modelforge.potential.parameters import (
+    ANI2xParameters,
+    PhysNetParameters,
+    SchNetParameters,
+    PaiNNParameters,
+    SAKEParameters,
+    TensorNetParameters,
+)
+from modelforge.train.parameters import TrainingParameters, RuntimeParameters
+from modelforge.dataset.dataset import DatasetParameters
+
+
 class NeuralNetworkPotentialFactory:
     """
     Factory class for creating instances of neural network potentials for training/inference.
     """
 
     @staticmethod
-    def generate_model(
+    def generate_potential(
         *,
         use: Literal["training", "inference"],
-        model_parameter: Dict[str, Union[str, Any]],
-        simulation_environment: Literal["PyTorch", "JAX"] = "PyTorch",
-        training_parameter: Optional[Dict[str, Any]] = None,
+        potential_parameter: Union[
+            ANI2xParameters,
+            SAKEParameters,
+            SchNetParameters,
+            PhysNetParameters,
+            PaiNNParameters,
+            TensorNetParameters,
+        ],
+        runtime_parameter: Optional[RuntimeParameters] = None,
+        training_parameter: Optional[TrainingParameters] = None,
+        dataset_parameter: Optional[DatasetParameters] = None,
         dataset_statistic: Optional[Dict[str, float]] = None,
+        potential_seed: Optional[int] = None,
+        simulation_environment: Optional[Literal["PyTorch", "JAX"]] = None,
     ) -> Union[Type[torch.nn.Module], Type[JAXModel], Type[pl.LightningModule]]:
         """
         Creates an NNP instance of the specified type, configured either for training or inference.
@@ -561,14 +583,21 @@ def generate_model(
         ----------
         use : Literal["training", "inference"]
             The use case for the model instance, either 'training' or 'inference'.
-        simulation_environment : Literal["PyTorch", "JAX"]
-            The ML framework to use, either 'PyTorch' or 'JAX'.
-        model_parameter : Dict[str, Union[str, Any]]
-            Parameters specific to the model.
-        training_parameter : Optional[Dict[str, Any]], optional
+        potential_parameter : Union[ANI2xParameters, SAKEParameters, SchNetParameters, PhysNetParameters, PaiNNParameters, TensorNetParameters]
+            Parameters specific to the potential.
+        runtime_parameter : Optional[RuntimeParameters], optional
+            Parameters for configuring the runtime environment.
+        training_parameter : Optional[TrainingParameters], optional
             Parameters for configuring the training.
+        dataset_parameter : Optional[DatasetParameters], optional
+            Parameters for configuring the dataset.
         dataset_statistic : Optional[Dict[str, float]], optional
             Statistics of the dataset for normalization purposes.
+        potential_seed : Optional[int], optional
+            Seed for the random number generator.
+        simulation_environment : Optional[Literal["PyTorch", "JAX"]], optional, None
+            The simulation environment to use for training/inference. Will override the runtime parameter if provided.
+
 
         Returns
         -------
@@ -584,33 +613,44 @@ def generate_model(
         """
 
         from modelforge.potential import _Implemented_NNPs
-        from modelforge.train.training import TrainingAdapter
+        from modelforge.train.training import ModelTrainer
 
         log.debug(f"{training_parameter=}")
-        log.debug(f"{model_parameter=}")
+        log.debug(f"{potential_parameter=}")
+        log.debug(f"{dataset_parameter=}")
+
+        if simulation_environment is None:
+            if runtime_parameter is None:
+                log.warning(
+                    "No runtime paramters or simulation_environment specified, defaulting to PyTorch"
+                )
 
+                simulation_environment = "PyTorch"
+            else:
+                simulation_environment = runtime_parameter.simulation_environment
         # obtain model for training
         if use == "training":
             if simulation_environment == "JAX":
                 log.warning(
                     "Training in JAX is not available. Falling back to PyTorch."
                 )
-            model = TrainingAdapter(
-                model_parameter=model_parameter,
-                lr_scheduler=training_parameter["lr_scheduler"],
-                lr=training_parameter["lr"],
-                loss_parameter=training_parameter["loss_parameter"],
-                dataset_statistic=dataset_statistic,
+            model = ModelTrainer(
+                potential_parameter=potential_parameter,
+                training_parameter=training_parameter,
+                dataset_parameter=dataset_parameter,
+                runtime_parameter=runtime_parameter,
+                potential_seed=potential_seed,
             )
             return model
         # obtain model for inference
         elif use == "inference":
-            model_type = model_parameter["potential_name"]
+            model_type = potential_parameter.potential_name
             nnp_class: Type = _Implemented_NNPs.get_neural_network_class(model_type)
             model = nnp_class(
-                **model_parameter["core_parameter"],
-                postprocessing_parameter=model_parameter["postprocessing_parameter"],
+                **potential_parameter.core_parameter.model_dump(),
+                postprocessing_parameter=potential_parameter.postprocessing_parameter.model_dump(),
                 dataset_statistic=dataset_statistic,
+                potential_seed=potential_seed,
             )
             if simulation_environment == "JAX":
                 return PyTorch2JAXConverter().convert_to_jax_model(model)
@@ -954,6 +994,7 @@ def __init__(
         postprocessing_parameter: Dict[str, Dict[str, bool]],
         dataset_statistic: Optional[Dict[str, float]],
         maximum_interaction_radius: unit.Quantity,
+        potential_seed: Optional[int] = None,
     ):
         """
         Initialize the BaseNetwork.
@@ -966,11 +1007,30 @@ def __init__(
             Dataset statistics for normalization.
         maximum_interaction_radius : unit.Quantity
             cutoff radius.
+        potential_seed : Optional[int], optional
+            Value used for torch.manual_seed, by default None.
         """
 
         super().__init__()
         from modelforge.utils.units import _convert_str_to_unit
 
+        if potential_seed:
+            import torch
+
+            torch.manual_seed(potential_seed)
+
+            # according to https://docs.ray.io/en/latest/tune/faq.html#how-can-i-reproduce-experiments
+            # we should also set the same seed for numpy.random and python random module
+            # when doing hyperparameter optimization with ray tune.  E.g., the ASHA scheduler relies on numpy.random
+            # and doesn't take a seed as an argument, so we need to set it here.
+            import numpy as np
+
+            np.random.seed(potential_seed)
+
+            import random
+
+            random.seed(potential_seed)
+
         self.postprocessing = PostProcessing(
             postprocessing_parameter, dataset_statistic
         )
@@ -1012,7 +1072,7 @@ def load_state_dict(
         """
 
         # Prefix to remove
-        prefix = "model."
+        prefix = "potential."
         excluded_keys = ["loss.per_molecule_energy", "loss.per_atom_force"]
 
         # Create a new dictionary without the prefix in the keys if prefix exists

modelforge/potential/painn.py

@@ -560,6 +560,7 @@ def __init__(
         postprocessing_parameter: Dict[str, Dict[str, bool]],
         dataset_statistic: Optional[Dict[str, float]] = None,
         epsilon: float = 1e-8,
+        potential_seed: Optional[int] = None,
     ) -> None:
         """
         Initialize the PaiNN network.
@@ -582,7 +583,8 @@ def __init__(
         shared_filters : bool
             Whether to share filters across modules.
             epsilon=epsilon,
-        )
+        potential_seed : Optional[int], optional
+            Seed for the random number generator, default None.
         """
 
         from modelforge.utils.units import _convert_str_to_unit
@@ -593,6 +595,7 @@ def __init__(
             dataset_statistic=dataset_statistic,
             postprocessing_parameter=postprocessing_parameter,
             maximum_interaction_radius=_convert_str_to_unit(maximum_interaction_radius),
+            potential_seed=potential_seed,
         )
 
         activation_function = activation_function_parameter["activation_function"]

modelforge/potential/parameters.py

@@ -153,6 +153,7 @@ class PostProcessingParameter(ParametersBase):
     potential_name: str = "ANI2x"
     core_parameter: CoreParameter
     postprocessing_parameter: PostProcessingParameter
+    potential_seed: Optional[int] = None
 
 
 class SchNetParameters(ParametersBase):
@@ -183,6 +184,7 @@ class PostProcessingParameter(ParametersBase):
     potential_name: str = "SchNet"
     core_parameter: CoreParameter
     postprocessing_parameter: PostProcessingParameter
+    potential_seed: Optional[int] = None
 
 
 class TensorNetParameters(ParametersBase):
@@ -214,6 +216,7 @@ class PostProcessingParameter(ParametersBase):
     potential_name: str = "TensorNet"
     core_parameter: CoreParameter
     postprocessing_parameter: PostProcessingParameter
+    potential_seed: Optional[int] = None
 
 
 class PaiNNParameters(ParametersBase):
@@ -244,6 +247,7 @@ class PostProcessingParameter(ParametersBase):
     potential_name: str = "PaiNN"
     core_parameter: CoreParameter
     postprocessing_parameter: PostProcessingParameter
+    potential_seed: Optional[int] = None
 
 
 class PhysNetParameters(ParametersBase):
@@ -273,6 +277,7 @@ class PostProcessingParameter(ParametersBase):
     potential_name: str = "PhysNet"
     core_parameter: CoreParameter
     postprocessing_parameter: PostProcessingParameter
+    potential_seed: Optional[int] = None
 
 
 class SAKEParameters(ParametersBase):
@@ -302,3 +307,4 @@ class PostProcessingParameter(ParametersBase):
     potential_name: str = "SAKE"
     core_parameter: CoreParameter
     postprocessing_parameter: PostProcessingParameter
+    potential_seed: Optional[int] = None

modelforge/potential/physnet.py

@@ -675,6 +675,8 @@ class PhysNet(BaseNetwork):
         Configuration for postprocessing parameters.
     dataset_statistic : Optional[Dict[str, float]], optional
         Statistics of the dataset, by default None.
+    potential_seed : Optional[int], optional
+        Seed for the random number generator, default None.
     """
 
     def __init__(
@@ -687,13 +689,15 @@ def __init__(
         activation_function_parameter: Dict,
         postprocessing_parameter: Dict[str, Dict[str, bool]],
         dataset_statistic: Optional[Dict[str, float]] = None,
+        potential_seed: Optional[int] = None,
     ) -> None:
 
         self.only_unique_pairs = False  # NOTE: for pairlist
         super().__init__(
             dataset_statistic=dataset_statistic,
             postprocessing_parameter=postprocessing_parameter,
             maximum_interaction_radius=_convert_str_to_unit(maximum_interaction_radius),
+            potential_seed=potential_seed,
         )
         activation_function = activation_function_parameter["activation_function"]
 

modelforge/potential/sake.py

@@ -602,6 +602,7 @@ def __init__(
         postprocessing_parameter: Dict[str, Dict[str, bool]],
         dataset_statistic: Optional[Dict[str, float]] = None,
         epsilon: float = 1e-8,
+        potential_seed: Optional[int] = None,
     ):
         from modelforge.utils.units import _convert_str_to_unit
 
@@ -610,6 +611,7 @@ def __init__(
             dataset_statistic=dataset_statistic,
             postprocessing_parameter=postprocessing_parameter,
             maximum_interaction_radius=_convert_str_to_unit(maximum_interaction_radius),
+            potential_seed=potential_seed,
         )
         activation_function = activation_function_parameter["activation_function"]
 

modelforge/potential/schnet.py

@@ -428,6 +428,8 @@ class SchNet(BaseNetwork):
         Configuration for postprocessing parameters.
     dataset_statistic : Optional[Dict[str, float]], default=None
         Statistics of the dataset.
+    potential_seed : Optional[int], optional
+        Seed for the random number generator, default None.
     """
 
     def __init__(
@@ -441,6 +443,7 @@ def __init__(
         shared_interactions: bool,
         postprocessing_parameter: Dict[str, Dict[str, bool]],
         dataset_statistic: Optional[Dict[str, float]] = None,
+        potential_seed: Optional[int] = None,
     ) -> None:
 
         self.only_unique_pairs = False  # NOTE: need to be set before super().__init__
@@ -449,6 +452,7 @@ def __init__(
             dataset_statistic=dataset_statistic,
             postprocessing_parameter=postprocessing_parameter,
             maximum_interaction_radius=_convert_str_to_unit(maximum_interaction_radius),
+            potential_seed=potential_seed,
         )
 
         activation_function = activation_function_parameter["activation_function"]

modelforge/potential/tensornet.py

@@ -203,6 +203,8 @@ class TensorNet(BaseNetwork):
         Postprocessing parameters.
     dataset_statistic : Optional[Dict[str, float]]
         Dataset statistics.
+    potential_seed : Optional[int], optional
+        Seed for the random number generator, default None.
     """
 
     def __init__(
@@ -217,6 +219,7 @@ def __init__(
         activation_function_parameter: Dict,
         postprocessing_parameter: Dict[str, Dict[str, bool]],
         dataset_statistic: Optional[Dict[str, float]] = None,
+        potential_seed: Optional[int] = None,
     ) -> None:
 
         activation_function = activation_function_parameter["activation_function"]
@@ -226,6 +229,7 @@ def __init__(
             dataset_statistic=dataset_statistic,
             postprocessing_parameter=postprocessing_parameter,
             maximum_interaction_radius=_convert_str_to_unit(maximum_interaction_radius),
+            potential_seed=potential_seed,
         )
 
         self.core_module = TensorNetCore(

modelforge/tests/data/config.toml

@@ -76,11 +76,11 @@ seed = 42
 
 [runtime]
 save_dir = "lightning_logs"
-experiment_name = "test_exp"
+experiment_name = "{potential_name}_{dataset_name}"
 local_cache_dir = "./cache"
 accelerator = "cpu"
 number_of_nodes = 1
 devices = 1                        #[0,1,2,3]
 checkpoint_path = "None"
 simulation_environment = "PyTorch"
-log_every_n_steps = 50
+log_every_n_steps = 1

modelforge/tests/data/runtime_defaults/runtime.toml

@@ -1,6 +1,6 @@
 [runtime]
 save_dir = "lightning_logs"
-experiment_name = "test_exp2"
+experiment_name = "{potential_name}_{dataset_name}"
 local_cache_dir = "./cache"
 accelerator = "cpu"
 number_of_nodes = 1