Repository: mlcommons/GaNDLF

Issue:
Add the ability to store models for every epoch
**Is your feature request related to a problem? Please describe.**
Currently, the model is getting saved based on validation loss, but depending on the desired metric, it might not be sufficient.

**Describe the solution you'd like**
Provide an option to the user to save models at **every single epoch**. When this option is selected, put a warning stating that this will blow the storage.

**Describe alternatives you've considered**
Provide a way to save models based on selected metric. This will get complicated unless we provide a way to map every metric to a common mathematical space (e.g., `0` is best and `1` is worst).

**Additional context**
Requested by @sbakas and @bahetibhakti 



Discussion:
Proposing to add a new option under `model` key, called `save_at_every_epoch`, which defaults to `False`.

Resolving patch:
diff --git a/GANDLF/compute/training_loop.py b/GANDLF/compute/training_loop.py
--- a/GANDLF/compute/training_loop.py
+++ b/GANDLF/compute/training_loop.py
@@ -492,6 +492,27 @@ def training_loop(
             model.train()
             first_model_saved = True
 
+        if params["model"]["save_at_every_epoch"]:
+            save_model(
+                {
+                    "epoch": epoch,
+                    "model_state_dict": model.state_dict(),
+                    "optimizer_state_dict": optimizer.state_dict(),
+                    "loss": epoch_valid_loss,
+                },
+                model,
+                params,
+                os.path.join(
+                    output_dir,
+                    params["model"]["architecture"]
+                    + "_epoch_"
+                    + str(epoch)
+                    + ".pth.tar",
+                ),
+                onnx_export=False,
+            )
+            model.train()
+
         print("Current Best epoch: ", best_train_idx)
 
         if patience > params["patience"]:
diff --git a/GANDLF/parseConfig.py b/GANDLF/parseConfig.py
--- a/GANDLF/parseConfig.py
+++ b/GANDLF/parseConfig.py
@@ -509,6 +509,15 @@ def parseConfig(config_file_path, version_check_flag=True):
         if not ("type" in params["model"]):
             params["model"]["type"] = "torch"
 
+        # set default save strategy for model
+        if not ("save_at_every_epoch" in params["model"]):
+            params["model"]["save_at_every_epoch"] = False
+
+        if params["model"]["save_at_every_epoch"]:
+            print(
+                "WARNING: 'save_at_every_epoch' will result in TREMENDOUS storage usage; use at your own risk."
+            )
+
     else:
         sys.exit("The 'model' parameter needs to be populated as a dictionary")
 


Test patch:
diff --git a/testing/test_full.py b/testing/test_full.py
--- a/testing/test_full.py
+++ b/testing/test_full.py
@@ -525,6 +525,7 @@ def test_train_resume_inference_classification_rad_3d(device):
     parameters["num_epochs"] = 2
     parameters["nested_training"]["testing"] = -5
     parameters["nested_training"]["validation"] = -5
+    parameters["model"]["save_at_every_epoch"] = True
     TrainingManager(
         dataframe=training_data,
         outputDir=outputDir,