Repository: mlcommons/GaNDLF

Issue:
Add checkpointing
**Is your feature request related to a problem? Please describe.**
Currently, GaNDLF is saving the "best" model automatically, but not loading it.

**Describe the solution you'd like**
During training, load the "best" model, if appropriate and update variables.

**Describe alternatives you've considered**
N.A.

**Additional context**
N.A.



Discussion:

Resolving patch:
diff --git a/GANDLF/training_loop.py b/GANDLF/training_loop.py
--- a/GANDLF/training_loop.py
+++ b/GANDLF/training_loop.py
@@ -821,14 +821,29 @@ def training_loop(
 
     # Setup a few variables for tracking
     best_loss = 1e7
-    patience = 0
+    patience, start_epoch = 0, 0
+    first_model_saved = False
+    best_model_path = os.path.join(
+        output_dir, params["model"]["architecture"] + "_best.pth.tar"
+    )
 
-    print("Using device:", device, flush=True)
+    # if previous model file is present, load it up
+    if os.path.exists(best_model_path):
+        print("Previous model found. Loading it up.")
+        try:
+            main_dict = torch.load(best_model_path)
+            model.load_state_dict(main_dict["model_state_dict"])
+            start_epoch = main_dict["epoch"]
+            optimizer.load_state_dict(main_dict["optimizer_state_dict"])
+            best_loss = main_dict["best_loss"]
+            print("Previous model loaded successfully.")
+        except Exception as e:
+            print("Previous model could not be loaded, error: ", e)
 
-    first_model_saved = False
+    print("Using device:", device, flush=True)
 
     # Iterate for number of epochs
-    for epoch in range(epochs):
+    for epoch in range(start_epoch, epochs):
 
         # Printing times
         epoch_start_time = time.time()
@@ -876,9 +891,7 @@ def training_loop(
                     "optimizer_state_dict": optimizer.state_dict(),
                     "best_loss": best_loss,
                 },
-                os.path.join(
-                    output_dir, params["model"]["architecture"] + "_best.pth.tar"
-                ),
+                best_model_path,
             )
             first_model_saved = True
 
diff --git a/GANDLF/training_manager.py b/GANDLF/training_manager.py
--- a/GANDLF/training_manager.py
+++ b/GANDLF/training_manager.py
@@ -21,21 +21,32 @@ def TrainingManager(dataframe, outputDir, parameters, device, reset_prev):
         shutil.rmtree(outputDir)
         Path(outputDir).mkdir(parents=True, exist_ok=True)
 
+    # save the current model configuration as a sanity check
+    currentModelConfigPickle = os.path.join(outputDir, "parameters.pkl")
+    if (not os.path.exists(currentModelConfigPickle)) or reset_prev:
+        with open(currentModelConfigPickle, "wb") as handle:
+            pickle.dump(parameters, handle, protocol=pickle.HIGHEST_PROTOCOL)
+    else:
+        print(
+            "Using previously saved parameter file",
+            currentModelConfigPickle,
+            flush=True,
+        )
+        parameters = pickle.load(open(currentModelConfigPickle, "rb"))
+
     # check for single fold training
     singleFoldValidation = False
     singleFoldTesting = False
     noTestingData = False
-    if (
-        parameters["nested_training"]["testing"] < 0
-    ):  # if the user wants a single fold training
+    # if the user wants a single fold training
+    if parameters["nested_training"]["testing"] < 0:
         parameters["nested_training"]["testing"] = abs(
             parameters["nested_training"]["testing"]
         )
         singleFoldTesting = True
 
-    if (
-        parameters["nested_training"]["validation"] < 0
-    ):  # if the user wants a single fold training
+    # if the user wants a single fold training
+    if parameters["nested_training"]["validation"] < 0:
         parameters["nested_training"]["validation"] = abs(
             parameters["nested_training"]["validation"]
         )
@@ -45,9 +56,8 @@ def TrainingManager(dataframe, outputDir, parameters, device, reset_prev):
     if parameters["nested_training"]["testing"] == 1:
         noTestingData = True
         singleFoldTesting = True
-        parameters["nested_training"][
-            "testing"
-        ] = 2  # put 2 just so that the first for-loop does not fail
+        # put 2 just so that the first for-loop does not fail
+        parameters["nested_training"]["testing"] = 2
 
     # initialize the kfold structures
     kf_testing = KFold(n_splits=parameters["nested_training"]["testing"])
@@ -66,9 +76,7 @@ def TrainingManager(dataframe, outputDir, parameters, device, reset_prev):
     trainingData_full = dataframe
 
     # start the kFold train for testing
-    for trainAndVal_index, testing_index in kf_testing.split(
-        subjectIDs_full
-    ):  # perform testing split
+    for trainAndVal_index, testing_index in kf_testing.split(subjectIDs_full):
 
         currentValidationFold = (
             0  # ensure the validation fold is initialized per-testing split
@@ -118,19 +126,6 @@ def TrainingManager(dataframe, outputDir, parameters, device, reset_prev):
             )
             Path(currentOutputFolder).mkdir(parents=True, exist_ok=True)
 
-        # save the current model configuration as a sanity check
-        currentModelConfigPickle = os.path.join(currentOutputFolder, "parameters.pkl")
-        if (not os.path.exists(currentModelConfigPickle)) or reset_prev:
-            with open(currentModelConfigPickle, "wb") as handle:
-                pickle.dump(parameters, handle, protocol=pickle.HIGHEST_PROTOCOL)
-        else:
-            print(
-                "Using previously saved parameter file",
-                currentModelConfigPickle,
-                flush=True,
-            )
-            parameters = pickle.load(open(currentModelConfigPickle, "rb"))
-
         # save the current training+validation and testing datasets
         if noTestingData:
             print("!!! WARNING !!!")


Test patch:
diff --git a/testing/test_full.py b/testing/test_full.py
--- a/testing/test_full.py
+++ b/testing/test_full.py
@@ -872,3 +872,45 @@ def test_augmentation_functions():
         output_tensor = global_augs_dict[aug]()(input_tensor)
         assert output_tensor != None, "Augmentation should work"
     print("passed")
+
+
+def test_checkpointing_segmentation_rad_2d(device):
+    print("Starting 2D Rad segmentation tests for metrics")
+    # read and parse csv
+    parameters = parseConfig(
+        testingDir + "/config_segmentation.yaml", version_check=False
+    )
+    training_data, parameters["headers"] = parseTrainingCSV(
+        inputDir + "/train_2d_rad_segmentation.csv"
+    )
+    parameters = populate_header_in_parameters(parameters, parameters["headers"])
+    parameters["patch_size"] = patch_size["2D"]
+    parameters["num_epochs"] = 1
+    parameters["nested_training"]["testing"] = 1
+    parameters["model"]["dimension"] = 2
+    parameters["model"]["class_list"] = [0, 255]
+    parameters["model"]["amp"] = True
+    parameters["model"]["num_channels"] = 3
+    parameters["metrics"] = ["dice", "hausdorff", "hausdorff95"]
+    parameters["model"]["architecture"] = "unet"
+    Path(outputDir).mkdir(parents=True, exist_ok=True)
+    TrainingManager(
+        dataframe=training_data,
+        outputDir=outputDir,
+        parameters=parameters,
+        device=device,
+        reset_prev=True,
+    )
+    parameters["num_epochs"] = 2
+    parameters["nested_training"]["validation"] = -2
+    parameters["nested_training"]["testing"] = 1
+    TrainingManager(
+        dataframe=training_data,
+        outputDir=outputDir,
+        parameters=parameters,
+        device=device,
+        reset_prev=False,
+    )
+    shutil.rmtree(outputDir)  # overwrite previous results
+
+    print("passed")