Repository: mlcommons/GaNDLF

Issue:
Ensure the training parameters are captured in the model dictionary itself
**Is your feature request related to a problem? Please describe.**
Currently, the training configuration needs to be passed separately to the model, which is cumbersome. 

**Describe the solution you'd like**
By directly including the parameters within the model dictionary, a potential breaking point can be fixed.

**Describe alternatives you've considered**
N.A.

**Additional context**
N.A.



Discussion:

Resolving patch:
diff --git a/GANDLF/cli/post_training_model_optimization.py b/GANDLF/cli/post_training_model_optimization.py
--- a/GANDLF/cli/post_training_model_optimization.py
+++ b/GANDLF/cli/post_training_model_optimization.py
@@ -16,7 +16,13 @@ def post_training_model_optimization(model_path, config_path):
         bool: True if successful, False otherwise.
     """
 
-    parameters = parseConfig(config_path, version_check_flag=False)
+    main_dict = load_model(model_path, "cpu")
+    parameters = main_dict.get("parameters", None)
+    parameters = (
+        parseConfig(config_path, version_check_flag=False)
+        if parameters is None
+        else parameters
+    )
     (
         model,
         _,
@@ -25,8 +31,8 @@ def post_training_model_optimization(model_path, config_path):
         _,
         parameters,
     ) = create_pytorch_objects(parameters, device="cpu")
+    parameters["model"]["onnx_export"] = True
 
-    main_dict = load_model(model_path, "cpu")
     version_check(parameters["version"], version_to_check=main_dict["version"])
     model.load_state_dict(main_dict["model_state_dict"])
     optimize_and_save_model(model, parameters, model_path, onnx_export=True)
diff --git a/GANDLF/utils/generic.py b/GANDLF/utils/generic.py
--- a/GANDLF/utils/generic.py
+++ b/GANDLF/utils/generic.py
@@ -94,7 +94,8 @@ def parse_version(version_string):
     Returns:
         int: The version number.
     """
-    version_string_split = version_string.split(".")
+    version_string_split = version_string.replace("-dev", "")
+    version_string_split = version_string_split.split(".")
     if len(version_string_split) > 3:
         del version_string_split[-1]
     return int("".join(version_string_split))
diff --git a/GANDLF/utils/modelio.py b/GANDLF/utils/modelio.py
--- a/GANDLF/utils/modelio.py
+++ b/GANDLF/utils/modelio.py
@@ -2,6 +2,7 @@
 import torch
 
 from .generic import get_unique_timestamp
+from ..version import __version__
 
 # these are the base keys for the model dictionary to save
 model_dict_full = {
@@ -150,7 +151,8 @@ def save_model(model_dict, model, params, path, onnx_export=True):
     model_dict["timestamp_hash"] = hashlib.sha256(
         str(model_dict["timestamp"]).encode("utf-8")
     ).hexdigest()
-    model_dict["version"] = pkg_resources.require("GANDLF")[0].version
+    model_dict["version"] = __version__
+    model_dict["parameters"] = params
     try:
         model_dict["git_hash"] = (
             subprocess.check_output(["git", "rev-parse", "HEAD"])


Test patch:
diff --git a/testing/test_full.py b/testing/test_full.py
--- a/testing/test_full.py
+++ b/testing/test_full.py
@@ -486,7 +486,8 @@ def test_train_regression_brainage_rad_2d(device):
     with open(file_config_temp, "w") as file:
         yaml.dump(parameters_temp, file)
     model_path = os.path.join(outputDir, "brain_age_best.pth.tar")
-    optimization_result = post_training_model_optimization(model_path, file_config_temp)
+    config_path = os.path.join(outputDir, "parameters.pkl")
+    optimization_result = post_training_model_optimization(model_path, config_path)
     assert optimization_result == False, "Optimization should fail"
 
     sanitize_outputDir()
@@ -759,7 +760,8 @@ def test_train_inference_optimize_classification_rad_3d(device):
     with open(file_config_temp, "w") as file:
         yaml.dump(parameters_temp, file)
     model_path = os.path.join(outputDir, all_models_regression[0] + "_best.pth.tar")
-    optimization_result = post_training_model_optimization(model_path, file_config_temp)
+    config_path = os.path.join(outputDir, "parameters.pkl")
+    optimization_result = post_training_model_optimization(model_path, config_path)
     assert optimization_result == True, "Optimization should pass"
 
     ## testing inference