Repository: mlcommons/GaNDLF

Issue:
Enable fine-tuning of previously trained models
**Is your feature request related to a problem? Please describe.**
Currently, if one wants to fine-tune previously trained models, they would need to do the following:
- Delete training data, validation data, and parameter pickles from `modeldir`
- Ensure the model.pth.tar file stays
- Resume training

**Describe the solution you'd like**
- Change the current `resetprev` to `reset`, and keep functionality the same (i.e., deletes **everything** from previous run)
- Add `resume`, which will only be used when `train == True`, and it only deletes the following pickles from `modeldir`: training data, validation data, and parameter
- If both `resume` and `reset` are set, assume `resume` and disable `reset` and provide a warning

**Describe alternatives you've considered**
N.A.

**Additional context**
N.A.



Discussion:

Resolving patch:
diff --git a/GANDLF/cli/main_run.py b/GANDLF/cli/main_run.py
--- a/GANDLF/cli/main_run.py
+++ b/GANDLF/cli/main_run.py
@@ -7,7 +7,7 @@
 from GANDLF.utils import populate_header_in_parameters, parseTrainingCSV
 
 
-def main_run(data_csv, config_file, output_dir, train_mode, device, reset_prev):
+def main_run(data_csv, config_file, output_dir, train_mode, device, resume, reset):
     """
     Main function that runs the training and inference.
 
@@ -17,7 +17,8 @@ def main_run(data_csv, config_file, output_dir, train_mode, device, reset_prev):
         output_dir (str): The output directory.
         train_mode (bool): Whether to train or infer.
         device (str): The device type.
-        reset_prev (bool): Whether the previous run will be reset or not.
+        resume (bool): Whether the previous run will be resumed or not.
+        reset (bool): Whether the previous run will be reset or not.
 
     Raises:
         ValueError: Parameter check from previous run.
@@ -28,10 +29,12 @@ def main_run(data_csv, config_file, output_dir, train_mode, device, reset_prev):
     parameters = parseConfig(model_parameters)
     # in case the data being passed is already processed, check if the previous parameters exists,
     # and if it does, compare the two and if they are the same, ensure no preprocess is done.
-    model_parameters_prev = os.path.join(
-        os.path.dirname(file_data_full), "parameters.pkl"
-    )
-    if not reset_prev:
+    model_parameters_prev = os.path.join(os.path.dirname(output_dir), "parameters.pkl")
+    if not (reset) or not (resume):
+        print(
+            "Trying to resume training without changing any parameters from previous run.",
+            flush=True,
+        )
         if os.path.exists(model_parameters_prev):
             parameters_prev = pickle.load(open(model_parameters_prev, "rb"))
             if parameters != parameters_prev:
@@ -39,11 +42,8 @@ def main_run(data_csv, config_file, output_dir, train_mode, device, reset_prev):
                     "The parameters are not the same as the ones stored in the previous run, please re-check."
                 )
 
-        parameters["data_preprocessing"] = {}
     parameters["output_dir"] = output_dir
 
-    reset_prev = reset_prev
-
     if "-1" in device:
         device = "cpu"
 
@@ -73,7 +73,8 @@ def main_run(data_csv, config_file, output_dir, train_mode, device, reset_prev):
             outputDir=parameters["output_dir"],
             parameters=parameters,
             device=device,
-            reset_prev=reset_prev,
+            resume=resume,
+            reset=reset,
         )
     else:
         data_full, headers = parseTrainingCSV(file_data_full, train=train_mode)
@@ -86,7 +87,8 @@ def main_run(data_csv, config_file, output_dir, train_mode, device, reset_prev):
             outputDir=parameters["output_dir"],
             parameters=parameters,
             device=device,
-            reset_prev=reset_prev,
+            resume=resume,
+            reset=reset,
         )
     else:
         InferenceManager(
diff --git a/GANDLF/training_manager.py b/GANDLF/training_manager.py
--- a/GANDLF/training_manager.py
+++ b/GANDLF/training_manager.py
@@ -7,32 +7,35 @@
 from GANDLF.compute import training_loop
 
 
-def TrainingManager(dataframe, outputDir, parameters, device, reset_prev):
+def TrainingManager(dataframe, outputDir, parameters, device, resume, reset):
     """
     This is the training manager that ties all the training functionality together
 
-    dataframe = full data from CSV
-    outputDir = the main output directory
-    parameters = parameters read from YAML
-    device = self-explanatory
-    reset_prev = whether the previous run in the same output directory is used or not
+    Args:
+        dataframe (pandas.DataFrame): The full data from CSV.
+        outputDir (str): The main output directory.
+        parameters (dict): The parameters dictionary.
+        device (str): The device to perform computations on.
+        resume (bool): Whether the previous run will be resumed or not.
+        reset (bool): Whether the previous run will be reset or not.
     """
-    if reset_prev:
+    if reset:
         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:
+    if (not os.path.exists(currentModelConfigPickle)) or reset or resume:
         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"))
+        if os.path.exists(currentModelConfigPickle):
+            print(
+                "Using previously saved parameter file",
+                currentModelConfigPickle,
+                flush=True,
+            )
+            parameters = pickle.load(open(currentModelConfigPickle, "rb"))
 
     # check for single fold training
     singleFoldValidation = False
@@ -139,14 +142,35 @@ def TrainingManager(dataframe, outputDir, parameters, device, reset_prev):
                 currentOutputFolder, "data_testing.pkl"
             )
 
-            if (not os.path.exists(currentTestingDataPickle)) or reset_prev:
+            if (not os.path.exists(currentTestingDataPickle)) or reset or resume:
                 testingData.to_pickle(currentTestingDataPickle)
+            else:
+                if os.path.exists(currentTestingDataPickle):
+                    print(
+                        "Using previously saved testing data",
+                        currentTestingDataPickle,
+                        flush=True,
+                    )
+                    testingData = pd.read_pickle(currentTestingDataPickle)
+
             if (
-                not os.path.exists(currentTrainingAndValidationDataPickle)
-            ) or reset_prev:
+                (not os.path.exists(currentTrainingAndValidationDataPickle))
+                or reset
+                or resume
+            ):
                 trainingAndValidationData.to_pickle(
                     currentTrainingAndValidationDataPickle
                 )
+            else:
+                if os.path.exists(currentTrainingAndValidationDataPickle):
+                    print(
+                        "Using previously saved training+validation data",
+                        currentTrainingAndValidationDataPickle,
+                        flush=True,
+                    )
+                    trainingAndValidationData = pd.read_pickle(
+                        currentTrainingAndValidationDataPickle
+                    )
 
             current_training_subject_indeces_full = (
                 trainingAndValidationData[
@@ -210,14 +234,26 @@ def TrainingManager(dataframe, outputDir, parameters, device, reset_prev):
             currentValidationDataPickle = os.path.join(
                 currentValOutputFolder, "data_validation.pkl"
             )
-            if (not os.path.exists(currentTrainingDataPickle)) or reset_prev:
+            if (not os.path.exists(currentTrainingDataPickle)) or reset or resume:
                 trainingData.to_pickle(currentTrainingDataPickle)
             else:
-                trainingData = pd.read_pickle(currentTrainingDataPickle)
-            if (not os.path.exists(currentValidationDataPickle)) or reset_prev:
+                if os.path.exists(currentTrainingDataPickle):
+                    print(
+                        "Using previously saved training data",
+                        currentTrainingDataPickle,
+                        flush=True,
+                    )
+                    trainingData = pd.read_pickle(currentTrainingDataPickle)
+            if (not os.path.exists(currentValidationDataPickle)) or reset or resume:
                 validationData.to_pickle(currentValidationDataPickle)
             else:
-                validationData = pd.read_pickle(currentValidationDataPickle)
+                if os.path.exists(currentValidationDataPickle):
+                    print(
+                        "Using previously saved validation data",
+                        currentValidationDataPickle,
+                        flush=True,
+                    )
+                    validationData = pd.read_pickle(currentValidationDataPickle)
 
             # parallel_compute_command is an empty string, thus no parallel computing requested
             if (not parameters["parallel_compute_command"]) or (singleFoldValidation):
@@ -275,29 +311,33 @@ def TrainingManager(dataframe, outputDir, parameters, device, reset_prev):
 
 
 def TrainingManager_split(
-    dataframe_train, dataframe_validation, outputDir, parameters, device, reset_prev
+    dataframe_train, dataframe_validation, outputDir, parameters, device, resume, reset
 ):
     """
     This is the training manager that ties all the training functionality together
 
-    dataframe_train = training data from CSV
-    dataframe_validation = validation data from CSV
-    outputDir = the main output directory
-    parameters = parameters read from YAML
-    device = self-explanatory
-    reset_prev = whether the previous run in the same output directory is used or not
+    Args:
+        dataframe_train (pandas.DataFrame): The training data from CSV.
+        dataframe_validation (pandas.DataFrame): The validation data from CSV.
+        outputDir (str): The main output directory.
+        parameters (dict): The parameters dictionary.
+        device (str): The device to perform computations on.
+        resume (bool): Whether the previous run will be resumed or not.
+        reset (bool): Whether the previous run will be reset or not.
     """
     currentModelConfigPickle = os.path.join(outputDir, "parameters.pkl")
-    if (not os.path.exists(currentModelConfigPickle)) or reset_prev:
+    if (not os.path.exists(currentModelConfigPickle)) or reset or resume:
         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"))
+        if os.path.exists(currentModelConfigPickle):
+            print(
+                "Using previously saved parameter file",
+                currentModelConfigPickle,
+                flush=True,
+            )
+            parameters = pickle.load(open(currentModelConfigPickle, "rb"))
+
     training_loop(
         training_data=dataframe_train,
         validation_data=dataframe_validation,


Test patch:
diff --git a/testing/test_full.py b/testing/test_full.py
--- a/testing/test_full.py
+++ b/testing/test_full.py
@@ -141,6 +141,12 @@ def test_constructTrainingCSV():
                 i += 1
 
 
+def sanitize_outputDir():
+    if os.path.isdir(outputDir):
+        shutil.rmtree(outputDir)  # overwrite previous results
+    Path(outputDir).mkdir(parents=True, exist_ok=True)
+
+
 def test_train_segmentation_rad_2d(device):
     print("Starting 2D Rad segmentation tests")
     # read and parse csv
@@ -161,14 +167,14 @@ def test_train_segmentation_rad_2d(device):
         parameters["model"]["architecture"] = model
         parameters["nested_training"]["testing"] = -5
         parameters["nested_training"]["validation"] = -5
-        shutil.rmtree(outputDir)  # overwrite previous results
-        Path(outputDir).mkdir(parents=True, exist_ok=True)
+        sanitize_outputDir()
         TrainingManager(
             dataframe=training_data,
             outputDir=outputDir,
             parameters=parameters,
             device=device,
-            reset_prev=True,
+            resume=False,
+            reset=True,
         )
 
     print("passed")
@@ -191,14 +197,14 @@ def test_train_segmentation_sdnet_rad_2d(device):
     parameters["model"]["num_channels"] = 1
     parameters["model"]["architecture"] = "sdnet"
     parameters = populate_header_in_parameters(parameters, parameters["headers"])
-    shutil.rmtree(outputDir)  # overwrite previous results
-    Path(outputDir).mkdir(parents=True, exist_ok=True)
+    sanitize_outputDir()
     TrainingManager(
         dataframe=training_data,
         outputDir=outputDir,
         parameters=parameters,
         device=device,
-        reset_prev=True,
+        resume=False,
+        reset=True,
     )
 
     print("passed")
@@ -226,14 +232,14 @@ def test_train_segmentation_rad_3d(device):
         parameters["model"]["architecture"] = model
         parameters["nested_training"]["testing"] = -5
         parameters["nested_training"]["validation"] = -5
-        shutil.rmtree(outputDir)  # overwrite previous results
-        Path(outputDir).mkdir(parents=True, exist_ok=True)
+        sanitize_outputDir()
         TrainingManager(
             dataframe=training_data,
             outputDir=outputDir,
             parameters=parameters,
             device=device,
-            reset_prev=True,
+            resume=False,
+            reset=True,
         )
 
     print("passed")
@@ -260,14 +266,14 @@ def test_train_regression_rad_2d(device):
         parameters["model"]["architecture"] = model
         parameters["nested_training"]["testing"] = -5
         parameters["nested_training"]["validation"] = -5
-        shutil.rmtree(outputDir)  # overwrite previous results
-        Path(outputDir).mkdir(parents=True, exist_ok=True)
+        sanitize_outputDir()
         TrainingManager(
             dataframe=training_data,
             outputDir=outputDir,
             parameters=parameters,
             device=device,
-            reset_prev=True,
+            resume=False,
+            reset=True,
         )
 
     print("passed")
@@ -290,14 +296,14 @@ def test_train_brainage_rad_2d(device):
     parameters["scaling_factor"] = 1
     parameters["model"]["architecture"] = "brain_age"
     parameters = populate_header_in_parameters(parameters, parameters["headers"])
-    shutil.rmtree(outputDir)  # overwrite previous results
-    Path(outputDir).mkdir(parents=True, exist_ok=True)
+    sanitize_outputDir()
     TrainingManager(
         dataframe=training_data,
         outputDir=outputDir,
         parameters=parameters,
         device=device,
-        reset_prev=True,
+        resume=False,
+        reset=True,
     )
 
     print("passed")
@@ -322,14 +328,14 @@ def test_train_regression_rad_3d(device):
         parameters["model"]["architecture"] = model
         parameters["nested_training"]["testing"] = -5
         parameters["nested_training"]["validation"] = -5
-        shutil.rmtree(outputDir)  # overwrite previous results
-        Path(outputDir).mkdir(parents=True, exist_ok=True)
+        sanitize_outputDir()
         TrainingManager(
             dataframe=training_data,
             outputDir=outputDir,
             parameters=parameters,
             device=device,
-            reset_prev=True,
+            resume=False,
+            reset=True,
         )
 
     print("passed")
@@ -355,14 +361,14 @@ def test_train_classification_rad_2d(device):
         parameters["model"]["architecture"] = model
         parameters["nested_training"]["testing"] = -5
         parameters["nested_training"]["validation"] = -5
-        shutil.rmtree(outputDir)  # overwrite previous results
-        Path(outputDir).mkdir(parents=True, exist_ok=True)
+        sanitize_outputDir()
         TrainingManager(
             dataframe=training_data,
             outputDir=outputDir,
             parameters=parameters,
             device=device,
-            reset_prev=True,
+            resume=False,
+            reset=True,
         )
 
     print("passed")
@@ -387,14 +393,14 @@ def test_train_classification_rad_3d(device):
         parameters["model"]["architecture"] = model
         parameters["nested_training"]["testing"] = -5
         parameters["nested_training"]["validation"] = -5
-        shutil.rmtree(outputDir)  # overwrite previous results
-        Path(outputDir).mkdir(parents=True, exist_ok=True)
+        sanitize_outputDir()
         TrainingManager(
             dataframe=training_data,
             outputDir=outputDir,
             parameters=parameters,
             device=device,
-            reset_prev=True,
+            resume=False,
+            reset=True,
         )
 
     print("passed")
@@ -417,15 +423,41 @@ def test_inference_classification_rad_3d(device):
     # loop through selected models and train for single epoch
     model = all_models_regression[0]
     parameters["model"]["architecture"] = model
-    Path(outputDir).mkdir(parents=True, exist_ok=True)
+    sanitize_outputDir()
     TrainingManager(
         dataframe=training_data,
         outputDir=outputDir,
         parameters=parameters,
         device=device,
-        reset_prev=True,
+        resume=False,
+        reset=True,
     )
-    parameters["output_dir"] = outputDir  # this is in inference mode
+    ## testing resume with parameter updates
+    parameters["num_epochs"] = 2
+    parameters["nested_training"]["testing"] = -5
+    parameters["nested_training"]["validation"] = -5
+    TrainingManager(
+        dataframe=training_data,
+        outputDir=outputDir,
+        parameters=parameters,
+        device=device,
+        resume=True,
+        reset=False,
+    )
+
+    ## testing resume without parameter updates
+    parameters["num_epochs"] = 3
+    TrainingManager(
+        dataframe=training_data,
+        outputDir=outputDir,
+        parameters=parameters,
+        device=device,
+        resume=False,
+        reset=False,
+    )
+
+    ## testing inference
+    parameters["output_dir"] = outputDir
     InferenceManager(
         dataframe=training_data,
         outputDir=outputDir,
@@ -455,13 +487,14 @@ def test_inference_classification_with_logits_single_fold_rad_3d(device):
     # loop through selected models and train for single epoch
     model = all_models_regression[0]
     parameters["model"]["architecture"] = model
-    Path(outputDir).mkdir(parents=True, exist_ok=True)
+    sanitize_outputDir()
     TrainingManager(
         dataframe=training_data,
         outputDir=outputDir,
         parameters=parameters,
         device=device,
-        reset_prev=True,
+        resume=False,
+        reset=True,
     )
     parameters["output_dir"] = outputDir  # this is in inference mode
     InferenceManager(
@@ -495,13 +528,14 @@ def test_inference_classification_with_logits_multiple_folds_rad_3d(device):
     # loop through selected models and train for single epoch
     model = all_models_regression[0]
     parameters["model"]["architecture"] = model
-    Path(outputDir).mkdir(parents=True, exist_ok=True)
+    sanitize_outputDir()
     TrainingManager(
         dataframe=training_data,
         outputDir=outputDir,
         parameters=parameters,
         device=device,
-        reset_prev=True,
+        resume=False,
+        reset=True,
     )
     parameters["output_dir"] = outputDir  # this is in inference mode
     InferenceManager(
@@ -544,10 +578,10 @@ def test_scheduler_classification_rad_2d(device):
             outputDir=outputDir,
             parameters=parameters,
             device=device,
-            reset_prev=True,
+            resume=False,
+            reset=True,
         )
 
-    shutil.rmtree(outputDir)
     print("passed")
 
 
@@ -581,10 +615,10 @@ def test_optimizer_classification_rad_2d(device):
             outputDir=outputDir,
             parameters=parameters,
             device=device,
-            reset_prev=True,
+            resume=False,
+            reset=True,
         )
 
-    shutil.rmtree(outputDir)
     print("passed")
 
 
@@ -609,16 +643,15 @@ def test_clip_train_classification_rad_3d(device):
         parameters["clip_mode"] = clip_mode
         parameters["nested_training"]["testing"] = -5
         parameters["nested_training"]["validation"] = -5
-        # shutil.rmtree(outputDir)  # overwrite previous results
-        Path(outputDir).mkdir(parents=True, exist_ok=True)
+        sanitize_outputDir()
         TrainingManager(
             dataframe=training_data,
             outputDir=outputDir,
             parameters=parameters,
             device=device,
-            reset_prev=True,
+            resume=False,
+            reset=True,
         )
-    shutil.rmtree(outputDir)  # overwrite previous results
     print("passed")
 
 
@@ -649,15 +682,18 @@ def test_normtype_train_segmentation_rad_3d(device):
             parameters["model"]["norm_type"] = norm
             parameters["nested_training"]["testing"] = -5
             parameters["nested_training"]["validation"] = -5
+            if os.path.isdir(outputDir):
+                shutil.rmtree(outputDir)  # overwrite previous results
             Path(outputDir).mkdir(parents=True, exist_ok=True)
             TrainingManager(
                 dataframe=training_data,
                 outputDir=outputDir,
                 parameters=parameters,
                 device=device,
-                reset_prev=True,
+                resume=False,
+                reset=True,
             )
-            shutil.rmtree(outputDir)  # overwrite previous results
+
         print("passed")
 
 
@@ -679,15 +715,15 @@ def test_metrics_segmentation_rad_2d(device):
     parameters["metrics"] = ["dice", "hausdorff", "hausdorff95"]
     parameters["model"]["architecture"] = "resunet"
     parameters = populate_header_in_parameters(parameters, parameters["headers"])
-    Path(outputDir).mkdir(parents=True, exist_ok=True)
+    sanitize_outputDir()
     TrainingManager(
         dataframe=training_data,
         outputDir=outputDir,
         parameters=parameters,
         device=device,
-        reset_prev=True,
+        resume=False,
+        reset=True,
     )
-    shutil.rmtree(outputDir)  # overwrite previous results
 
     print("passed")
 
@@ -709,15 +745,15 @@ def test_metrics_regression_rad_2d(device):
     parameters["model"]["num_channels"] = 3
     parameters["model"]["architecture"] = "vgg11"
     parameters = populate_header_in_parameters(parameters, parameters["headers"])
-    Path(outputDir).mkdir(parents=True, exist_ok=True)
+    sanitize_outputDir()
     TrainingManager(
         dataframe=training_data,
         outputDir=outputDir,
         parameters=parameters,
         device=device,
-        reset_prev=True,
+        resume=False,
+        reset=True,
     )
-    shutil.rmtree(outputDir)  # overwrite previous results
 
     print("passed")
 
@@ -745,22 +781,23 @@ def test_losses_segmentation_rad_2d(device):
         parameters["loss_function"] = loss_type
         parameters["nested_training"]["testing"] = -5
         parameters["nested_training"]["validation"] = -5
-        Path(outputDir).mkdir(parents=True, exist_ok=True)
+        sanitize_outputDir()
         TrainingManager(
             dataframe=training_data,
             outputDir=outputDir,
             parameters=parameters,
             device=device,
-            reset_prev=True,
+            resume=False,
+            reset=True,
         )
-        shutil.rmtree(outputDir)  # overwrite previous results
+
     print("passed")
 
 
 def test_config_read():
     print("Starting testing reading configuration")
     # read and parse csv
-    file_config_temp = os.path.join(testingDir, "config_segmentation_temp.yaml")
+    file_config_temp = os.path.join(outputDir, "config_segmentation_temp.yaml")
     # if found in previous run, discard.
     if os.path.exists(file_config_temp):
         os.remove(file_config_temp)
@@ -854,13 +891,8 @@ def test_config_read():
 def test_cli_function_preprocess():
     print("Starting testing cli function preprocess")
     file_config = os.path.join(testingDir, "config_segmentation.yaml")
-    file_config_temp = os.path.join(testingDir, "config_segmentation_temp.yaml")
-    # if found in previous run, discard.
-    if os.path.exists(file_config_temp):
-        os.remove(file_config_temp)
-        parameter_pickle_file = os.path.join(outputDir, "parameters.pkl")
-        if os.path.exists(parameter_pickle_file):
-            os.remove(parameter_pickle_file)
+    sanitize_outputDir()
+    file_config_temp = os.path.join(outputDir, "config_segmentation_temp.yaml")
     file_data = os.path.join(inputDir, "train_2d_rad_segmentation.csv")
 
     parameters = parseConfig(file_config)
@@ -889,7 +921,6 @@ def test_cli_function_preprocess():
     # check that the length of training data is what we expect
     assert len(training_data) == 10, "Number of rows in dataframe is not 10"
 
-    shutil.rmtree(outputDir)  # overwrite previous results
     print("passed")
 
 
@@ -898,7 +929,7 @@ def test_cli_function_mainrun(device):
     parameters = parseConfig(
         testingDir + "/config_segmentation.yaml", version_check_flag=False
     )
-    file_config_temp = os.path.join(testingDir, "config_segmentation_temp.yaml")
+    file_config_temp = os.path.join(outputDir, "config_segmentation_temp.yaml")
     # if found in previous run, discard.
     if os.path.exists(file_config_temp):
         os.remove(file_config_temp)
@@ -920,8 +951,11 @@ def test_cli_function_mainrun(device):
 
     file_data = os.path.join(inputDir, "train_2d_rad_segmentation.csv")
 
-    main_run(file_data, file_config_temp, outputDir, True, device, True)
-    shutil.rmtree(outputDir)  # overwrite previous results
+    main_run(
+        file_data, file_config_temp, outputDir, True, device, resume=False, reset=True
+    )
+    if os.path.isdir(outputDir):
+        shutil.rmtree(outputDir)  # overwrite previous results
     print("passed")
 
 
@@ -962,15 +996,16 @@ def test_dataloader_construction_train_segmentation_3d(device):
     parameters["weighted_loss"] = False
     parameters = populate_header_in_parameters(parameters, parameters["headers"])
     # loop through selected models and train for single epoch
-    Path(outputDir).mkdir(parents=True, exist_ok=True)
+    sanitize_outputDir()
     TrainingManager(
         dataframe=training_data,
         outputDir=outputDir,
         parameters=parameters,
         device=device,
-        reset_prev=True,
+        resume=False,
+        reset=True,
     )
-    shutil.rmtree(outputDir)  # overwrite previous results
+
     print("passed")
 
 
@@ -1113,13 +1148,14 @@ def test_checkpointing_segmentation_rad_2d(device):
     ]
     parameters["model"]["architecture"] = "unet"
     parameters = populate_header_in_parameters(parameters, parameters["headers"])
-    Path(outputDir).mkdir(parents=True, exist_ok=True)
+    sanitize_outputDir()
     TrainingManager(
         dataframe=training_data,
         outputDir=outputDir,
         parameters=parameters,
         device=device,
-        reset_prev=True,
+        resume=False,
+        reset=True,
     )
     parameters["num_epochs"] = 2
     parameters["nested_training"]["validation"] = -2
@@ -1129,9 +1165,9 @@ def test_checkpointing_segmentation_rad_2d(device):
         outputDir=outputDir,
         parameters=parameters,
         device=device,
-        reset_prev=False,
+        resume=False,
+        reset=False,
     )
-    shutil.rmtree(outputDir)  # overwrite previous results
 
     print("passed")
 
@@ -1251,7 +1287,9 @@ def test_anonymizer():
 
 def test_train_inference_segmentation_histology_2d(device):
     print("Starting histology train/inference tests")
-    output_dir_patches = os.path.join(testingDir, "histo_patches")
+    # overwrite previous results
+    sanitize_outputDir()
+    output_dir_patches = os.path.join(outputDir, "histo_patches")
     if os.path.isdir(output_dir_patches):
         shutil.rmtree(output_dir_patches)
     Path(output_dir_patches).mkdir(parents=True, exist_ok=True)
@@ -1295,29 +1333,27 @@ def test_train_inference_segmentation_histology_2d(device):
     parameters["nested_training"]["testing"] = 1
     parameters["nested_training"]["validation"] = -2
     parameters["metrics"] = ["dice"]
-    # overwrite previous results
-    if os.path.isdir(outputDir):
-        shutil.rmtree(outputDir)
-    Path(outputDir).mkdir(parents=True, exist_ok=True)
+    modelDir = os.path.join(outputDir, "modelDir")
+    if os.path.isdir(modelDir):
+        shutil.rmtree(modelDir)
+    Path(modelDir).mkdir(parents=True, exist_ok=True)
     TrainingManager(
         dataframe=training_data,
-        outputDir=outputDir,
+        outputDir=modelDir,
         parameters=parameters,
         device=device,
-        reset_prev=True,
+        resume=False,
+        reset=True,
     )
-    parameters["output_dir"] = outputDir  # this is in inference mode
+    parameters["output_dir"] = modelDir  # this is in inference mode
     inference_data, parameters["headers"] = parseTrainingCSV(
         inputDir + "/train_2d_histo_segmentation.csv", train=False
     )
-
     InferenceManager(
         dataframe=inference_data,
-        outputDir=outputDir,
+        outputDir=modelDir,
         parameters=parameters,
         device=device,
     )
 
-    shutil.rmtree(output_dir_patches)
-
     print("passed")