Split validation and test infer config (#502)

Split validation, test, ensemble inference flags

Author: JonathanTripp

CHANGELOG.md

@@ -13,10 +13,12 @@ created.
 ## Upcoming
 
 ### Added
+- ([#502](https://github.com/microsoft/InnerEye-DeepLearning/pull/502)) More flags for fine control of when to run inference.
 - ([#492](https://github.com/microsoft/InnerEye-DeepLearning/pull/492)) Adding capability for regression tests for test
 jobs that run in AzureML.
 
 ### Changed
+- ([#502](https://github.com/microsoft/InnerEye-DeepLearning/pull/502)) Renamed command line option 'perform_training_set_inference' to 'inference_on_train_set'. Replaced command line option 'perform_validation_and_test_set_inference' with the pair of options 'inference_on_val_set' and 'inference_on_test_set'.
 - ([#496](https://github.com/microsoft/InnerEye-DeepLearning/pull/496)) All plots are now saved as PNG, rather than JPG.
 - ([#497](https://github.com/microsoft/InnerEye-DeepLearning/pull/497)) Reducing the size of the code snapshot that
 gets uploaded to AzureML, by skipping all test folders.

InnerEye/Azure/azure_config.py

@@ -19,7 +19,7 @@
 from azureml.train.hyperdrive import HyperDriveConfig
 from git import Repo
 
-from InnerEye.Azure.azure_util import fetch_run, is_offline_run_context
+from InnerEye.Azure.azure_util import fetch_run, is_offline_run_context, remove_arg
 from InnerEye.Azure.secrets_handling import SecretsHandling, read_all_settings
 from InnerEye.Common import fixed_paths
 from InnerEye.Common.generic_parsing import GenericConfig
@@ -324,31 +324,7 @@ def set_script_params_except_submit_flag(self) -> None:
         Populates the script_param field of the present object from the arguments in sys.argv, with the exception
         of the "azureml" flag.
         """
-        args = sys.argv[1:]
-        submit_flag = f"--{AZURECONFIG_SUBMIT_TO_AZUREML}"
-        retained_args = []
-        i = 0
-        while i < len(args):
-            arg = args[i]
-            if arg.startswith(submit_flag):
-                if len(arg) == len(submit_flag):
-                    # The commandline argument is "--azureml", with something possibly following: This can either be
-                    # "--azureml True" or "--azureml --some_other_param"
-                    if i < (len(args) - 1):
-                        # If the next argument starts with a "-" then assume that it does not belong to the --azureml
-                        # flag. If there is no "-", assume it belongs to the --azureml flag, and skip both
-                        if not args[i + 1].startswith("-"):
-                            i = i + 1
-                elif arg[len(submit_flag)] == "=":
-                    # The commandline argument is "--azureml=True" or "--azureml=False": Continue with next arg
-                    pass
-                else:
-                    # The argument list contains a flag like "--azureml_foo": Keep that.
-                    retained_args.append(arg)
-            else:
-                retained_args.append(arg)
-            i = i + 1
-        self.script_params = retained_args
+        self.script_params = remove_arg(AZURECONFIG_SUBMIT_TO_AZUREML, sys.argv[1:])
 
 
 @dataclass

InnerEye/Azure/azure_util.py

@@ -455,3 +455,43 @@ def step_up_directories(path: Path) -> Generator[Path, None, None]:
         if parent == path:
             break
         path = parent
+
+
+def remove_arg(arg: str, args: List[str]) -> List[str]:
+    """
+    Remove an argument from a list of arguments. The argument list is assumed to contain
+    elements of the form:
+    "-a", "--arg1", "--arg2", "value2", or "--arg3=value"
+    If there is an item matching "--arg" then it will be removed from the list.
+
+    :param arg: Argument to look for.
+    :param args: List of arguments to scan.
+    :return: List of arguments with --arg removed, if present.
+    """
+    arg_opt = f"--{arg}"
+    no_arg_opt = f"--no-{arg}"
+    retained_args = []
+    i = 0
+    while i < len(args):
+        arg = args[i]
+        if arg.startswith(arg_opt):
+            if len(arg) == len(arg_opt):
+                # The commandline argument is "--arg", with something possibly following: This can either be
+                # "--arg_opt value" or "--arg_opt --some_other_param"
+                if i < (len(args) - 1):
+                    # If the next argument starts with a "-" then assume that it does not belong to the --arg
+                    # argument. If there is no "-", assume it belongs to the --arg_opt argument, and skip both
+                    if not args[i + 1].startswith("-"):
+                        i = i + 1
+            elif arg[len(arg_opt)] == "=":
+                # The commandline argument is "--arg=value": Continue with next arg
+                pass
+            else:
+                # The argument list contains an argument like "--arg_other_param": Keep that.
+                retained_args.append(arg)
+        elif arg == no_arg_opt:
+            pass
+        else:
+            retained_args.append(arg)
+        i = i + 1
+    return retained_args

InnerEye/ML/SSL/lightning_containers/ssl_container.py

@@ -121,8 +121,9 @@ def setup(self) -> None:
                                                            dataset_path=self.local_dataset,
                                                            batch_size=self.ssl_training_batch_size)})
         self.data_module: InnerEyeDataModuleTypes = self.get_data_module()
-        self.perform_validation_and_test_set_inference = False
-        if self.number_of_cross_validation_splits > 1:
+        self.inference_on_val_set = False
+        self.inference_on_test_set = False
+        if self.perform_cross_validation:
             raise NotImplementedError("Cross-validation logic is not implemented for this module.")
 
     def _load_config(self) -> None:

InnerEye/ML/deep_learning_config.py

@@ -5,17 +5,16 @@
 from __future__ import annotations
 
 import logging
-from enum import Enum, unique
-from pathlib import Path
-from typing import Any, Dict, List, Optional
-
 import param
+from enum import Enum, unique
 from pandas import DataFrame
 from param import Parameterized
+from pathlib import Path
+from typing import Any, Dict, List, Optional
 
 from InnerEye.Azure.azure_util import DEFAULT_CROSS_VALIDATION_SPLIT_INDEX, RUN_CONTEXT, is_offline_run_context
 from InnerEye.Common import fixed_paths
-from InnerEye.Common.common_util import is_windows
+from InnerEye.Common.common_util import ModelProcessing, is_windows
 from InnerEye.Common.fixed_paths import DEFAULT_AML_UPLOAD_DIR, DEFAULT_LOGS_DIR_NAME
 from InnerEye.Common.generic_parsing import GenericConfig
 from InnerEye.Common.type_annotations import PathOrString, TupleFloat2
@@ -199,14 +198,24 @@ class WorkflowParams(param.Parameterized):
     cross_validation_split_index: int = param.Integer(DEFAULT_CROSS_VALIDATION_SPLIT_INDEX, bounds=(-1, None),
                                                       doc="The index of the cross validation fold this model is "
                                                           "associated with when performing k-fold cross validation")
-    perform_training_set_inference: bool = \
-        param.Boolean(False,
-                      doc="If True, run full image inference on the training set at the end of training. If False and "
-                          "perform_validation_and_test_set_inference is True (default), only run inference on "
-                          "validation and test set. If both flags are False do not run inference.")
-    perform_validation_and_test_set_inference: bool = \
-        param.Boolean(True,
-                      doc="If True (default), run full image inference on validation and test set after training.")
+    inference_on_train_set: Optional[bool] = \
+        param.Boolean(None,
+                      doc="If set, enable/disable full image inference on training set after training.")
+    inference_on_val_set: Optional[bool] = \
+        param.Boolean(None,
+                      doc="If set, enable/disable full image inference on validation set after training.")
+    inference_on_test_set: Optional[bool] = \
+        param.Boolean(None,
+                      doc="If set, enable/disable full image inference on test set after training.")
+    ensemble_inference_on_train_set: Optional[bool] = \
+        param.Boolean(None,
+                      doc="If set, enable/disable full image inference on the training set after ensemble training.")
+    ensemble_inference_on_val_set: Optional[bool] = \
+        param.Boolean(None,
+                      doc="If set, enable/disable full image inference on validation set after ensemble training.")
+    ensemble_inference_on_test_set: Optional[bool] = \
+        param.Boolean(None,
+                      doc="If set, enable/disable full image inference on test set after ensemble training.")
     weights_url: str = param.String(doc="If provided, a url from which weights will be downloaded and used for model "
                                         "initialization.")
     local_weights_path: Optional[Path] = param.ClassSelector(class_=Path,
@@ -254,6 +263,53 @@ def validate(self) -> None:
                              f"found number_of_cross_validation_splits = {self.number_of_cross_validation_splits} "
                              f"and cross_validation_split_index={self.cross_validation_split_index}")
 
+    """ Defaults for when to run inference in the absence of any command line switches. """
+    INFERENCE_DEFAULTS: Dict[ModelProcessing, Dict[ModelExecutionMode, bool]] = {
+        ModelProcessing.DEFAULT: {
+            ModelExecutionMode.TRAIN: False,
+            ModelExecutionMode.TEST: True,
+            ModelExecutionMode.VAL: True,
+        },
+        ModelProcessing.ENSEMBLE_CREATION: {
+            ModelExecutionMode.TRAIN: False,
+            ModelExecutionMode.TEST: True,
+            ModelExecutionMode.VAL: False,
+        }
+    }
+
+    def inference_options(self) -> Dict[ModelProcessing, Dict[ModelExecutionMode, Optional[bool]]]:
+        """
+        Return a mapping from ModelProcesing and ModelExecutionMode to command line switch.
+
+        :return: Command line switch for each combination of ModelProcessing and ModelExecutionMode.
+        """
+        return {
+            ModelProcessing.DEFAULT: {
+                ModelExecutionMode.TRAIN: self.inference_on_train_set,
+                ModelExecutionMode.TEST: self.inference_on_test_set,
+                ModelExecutionMode.VAL: self.inference_on_val_set,
+            },
+            ModelProcessing.ENSEMBLE_CREATION: {
+                ModelExecutionMode.TRAIN: self.ensemble_inference_on_train_set,
+                ModelExecutionMode.TEST: self.ensemble_inference_on_test_set,
+                ModelExecutionMode.VAL: self.ensemble_inference_on_val_set,
+            }
+        }
+
+    def inference_on_set(self, model_proc: ModelProcessing, data_split: ModelExecutionMode) -> bool:
+        """
+        Returns True if inference is required for this model_proc and data_split.
+
+        :param model_proc: Whether we are testing an ensemble or single model.
+        :param data_split: Indicates which of the 3 sets (training, test, or validation) is being processed.
+        :return: True if inference required.
+        """
+        inference_option = self.inference_options()[model_proc][data_split]
+        if inference_option is not None:
+            return inference_option
+
+        return WorkflowParams.INFERENCE_DEFAULTS[model_proc][data_split]
+
     @property
     def is_offline_run(self) -> bool:
         """

InnerEye/ML/run_ml.py

@@ -11,6 +11,7 @@
 from typing import Any, Callable, Dict, List, Optional, Tuple
 
 import pandas as pd
+from pytorch_lightning.core.datamodule import LightningDataModule
 import stopit
 import torch.multiprocessing
 from azureml._restclient.constants import RunStatus
@@ -120,19 +121,16 @@ def download_dataset(azure_dataset_id: str,
     return expected_dataset_path
 
 
-def log_metrics(val_metrics: Optional[InferenceMetricsForSegmentation],
-                test_metrics: Optional[InferenceMetricsForSegmentation],
-                train_metrics: Optional[InferenceMetricsForSegmentation],
+def log_metrics(metrics: Dict[ModelExecutionMode, InferenceMetrics],
                 run_context: Run) -> None:
     """
     Log metrics for each split to the provided run, or the current run context if None provided
-    :param val_metrics: Inference results for the validation split
-    :param test_metrics: Inference results for the test split
-    :param train_metrics: Inference results for the train split
+    :param metrics: Dictionary of inference results for each split.
     :param run_context: Run for which to log the metrics to, use the current run context if None provided
     """
-    for split in [x for x in [val_metrics, test_metrics, train_metrics] if x]:
-        split.log_metrics(run_context)
+    for split in metrics.values():
+        if isinstance(split, InferenceMetricsForSegmentation):
+            split.log_metrics(run_context)
 
 
 class MLRunner:
@@ -390,7 +388,7 @@ def run(self) -> None:
 
                 # If this is an cross validation run, and the present run is child run 0, then wait for the sibling
                 # runs, build the ensemble model, and write a report for that.
-                if self.container.number_of_cross_validation_splits > 0:
+                if self.container.perform_cross_validation:
                     should_wait_for_other_child_runs = (not self.is_offline_run) and \
                                                        self.container.cross_validation_split_index == 0
                     if should_wait_for_other_child_runs:
@@ -420,10 +418,24 @@ def is_normal_run_or_crossval_child_0(self) -> bool:
         """
         Returns True if the present run is a non-crossvalidation run, or child run 0 of a crossvalidation run.
         """
-        if self.container.number_of_cross_validation_splits > 0:
+        if self.container.perform_cross_validation:
             return self.container.cross_validation_split_index == 0
         return True
 
+    @staticmethod
+    def lightning_data_module_dataloaders(data: LightningDataModule) -> Dict[ModelExecutionMode, Callable]:
+        """
+        Given a lightning data module, return a dictionary of dataloader for each model execution mode.
+
+        :param data: Lightning data module.
+        :return: Data loader for each model execution mode.
+        """
+        return {
+            ModelExecutionMode.TEST: data.test_dataloader,
+            ModelExecutionMode.VAL: data.val_dataloader,
+            ModelExecutionMode.TRAIN: data.train_dataloader
+        }
+
     def run_inference_for_lightning_models(self, checkpoint_paths: List[Path]) -> None:
         """
         Run inference on the test set for all models that are specified via a LightningContainer.
@@ -439,11 +451,10 @@ def run_inference_for_lightning_models(self, checkpoint_paths: List[Path]) -> No
             # Read the data modules before changing the working directory, in case the code relies on relative paths
             data = self.container.get_inference_data_module()
             dataloaders: List[Tuple[DataLoader, ModelExecutionMode]] = []
-            if self.container.perform_validation_and_test_set_inference:
-                dataloaders.append((data.test_dataloader(), ModelExecutionMode.TEST))  # type: ignore
-                dataloaders.append((data.val_dataloader(), ModelExecutionMode.VAL))  # type: ignore
-            if self.container.perform_training_set_inference:
-                dataloaders.append((data.train_dataloader(), ModelExecutionMode.TRAIN))  # type: ignore
+            data_dataloaders = MLRunner.lightning_data_module_dataloaders(data)
+            for data_split, dataloader in data_dataloaders.items():
+                if self.container.inference_on_set(ModelProcessing.DEFAULT, data_split):
+                    dataloaders.append((dataloader(), data_split))
             checkpoint = load_checkpoint(checkpoint_paths[0], use_gpu=self.container.use_gpu)
             lightning_model.load_state_dict(checkpoint['state_dict'])
             lightning_model.eval()
@@ -491,8 +502,8 @@ def run_inference(self, checkpoint_handler: CheckpointHandler,
         """
 
         # run full image inference on existing or newly trained model on the training, and testing set
-        test_metrics, val_metrics, _ = self.model_inference_train_and_test(checkpoint_handler=checkpoint_handler,
-                                                                           model_proc=model_proc)
+        self.model_inference_train_and_test(checkpoint_handler=checkpoint_handler,
+                                            model_proc=model_proc)
 
         self.try_compare_scores_against_baselines(model_proc)
 
@@ -752,37 +763,25 @@ def copy_file(source: Path, destination_file: str) -> None:
     def model_inference_train_and_test(self,
                                        checkpoint_handler: CheckpointHandler,
                                        model_proc: ModelProcessing = ModelProcessing.DEFAULT) -> \
-            Tuple[Optional[InferenceMetrics], Optional[InferenceMetrics], Optional[InferenceMetrics]]:
-        train_metrics = None
-        val_metrics = None
-        test_metrics = None
+            Dict[ModelExecutionMode, InferenceMetrics]:
+        metrics: Dict[ModelExecutionMode, InferenceMetrics] = {}
 
         config = self.innereye_config
 
-        def run_model_test(data_split: ModelExecutionMode) -> Optional[InferenceMetrics]:
-            return model_test(config, data_split=data_split, checkpoint_handler=checkpoint_handler,  # type: ignore
-                              model_proc=model_proc)
-
-        if config.perform_validation_and_test_set_inference:
-            # perform inference on test set
-            test_metrics = run_model_test(ModelExecutionMode.TEST)
-            # perform inference on validation set (not for ensemble as current val is in the training fold
-            # for at least one of the models).
-            if model_proc != ModelProcessing.ENSEMBLE_CREATION:
-                val_metrics = run_model_test(ModelExecutionMode.VAL)
-
-        if config.perform_training_set_inference:
-            # perform inference on training set if required
-            train_metrics = run_model_test(ModelExecutionMode.TRAIN)
+        for data_split in ModelExecutionMode:
+            if self.container.inference_on_set(model_proc, data_split):
+                opt_metrics = model_test(config, data_split=data_split, checkpoint_handler=checkpoint_handler,
+                                         model_proc=model_proc)
+                if opt_metrics is not None:
+                    metrics[data_split] = opt_metrics
 
         # log the metrics to AzureML experiment if possible. When doing ensemble runs, log to the Hyperdrive parent run,
         # so that we get the metrics of child run 0 and the ensemble separated.
         if config.is_segmentation_model and not self.is_offline_run:
             run_for_logging = PARENT_RUN_CONTEXT if model_proc.ENSEMBLE_CREATION else RUN_CONTEXT
-            log_metrics(val_metrics=val_metrics, test_metrics=test_metrics,  # type: ignore
-                        train_metrics=train_metrics, run_context=run_for_logging)  # type: ignore
+            log_metrics(metrics=metrics, run_context=run_for_logging)  # type: ignore
 
-        return test_metrics, val_metrics, train_metrics
+        return metrics
 
     @stopit.threading_timeoutable()
     def wait_for_runs_to_finish(self, delay: int = 60) -> None:

Tests/ML/configs/lightning_test_containers.py

@@ -190,7 +190,7 @@ class DummyContainerWithModel(LightningContainer):
 
     def __init__(self) -> None:
         super().__init__()
-        self.perform_training_set_inference = True
+        self.inference_on_train_set = True
         self.num_epochs = 50
         self.l_rate = 1e-1
 

Tests/ML/models/test_scalar_model.py

@@ -313,8 +313,9 @@ def test_run_ml_with_segmentation_model(test_output_dirs: OutputFolderForTests)
     # This is for a bug in an earlier version of the code where the wrong execution mode was used to
     # compute the expected mask size at training time.
     config.test_crop_size = (75, 75, 75)
-    config.perform_training_set_inference = False
-    config.perform_validation_and_test_set_inference = True
+    config.inference_on_train_set = False
+    config.inference_on_val_set = True
+    config.inference_on_test_set = True
     config.set_output_to(test_output_dirs.root_dir)
     azure_config = get_default_azure_config()
     azure_config.train = True