Repository: Project-MONAI/MONAI
File: monai/apps/nnunet/nnunetv2_runner.py (Python)

Code before:
        - 5-fold training for a single model on 2 GPUs. Here ``configs`` is used to specify the configurations.

        .. code-block:: bash

            python -m monai.apps.nnunet nnUNetV2Runner train --input_config "./input.yaml" \\
                --configs "3d_fullres" \\
                --trainer_class_name "nnUNetTrainer_5epochs" \\
                --export_validation_probabilities True \\
                --device_ids "(0,1)"

        - find the best configuration

        .. code-block:: bash

            python -m monai.apps.nnunet nnUNetV2Runner find_best_configuration --input_config "./input.yaml"

        - predict, ensemble, and post-process

        .. code-block:: bash

            python -m monai.apps.nnunet nnUNetV2Runner predict_ensemble_postprocessing --input_config "./input.yaml"

    """

    def __init__(
        self, input_config: Any, trainer_class_name: str = "nnUNetTrainer", work_dir: str = "work_dir"
    ) -> None:
        self.input_info: dict = {}
        self.input_config_or_dict = input_config
        self.trainer_class_name = trainer_class_name
        self.work_dir = work_dir

        if isinstance(self.input_config_or_dict, dict):
            self.input_info = self.input_config_or_dict
        elif isinstance(self.input_config_or_dict, str) and os.path.isfile(self.input_config_or_dict):
            self.input_info = ConfigParser.load_config_file(self.input_config_or_dict)
        else:
            raise ValueError(f"{input_config} is not a valid file or dict")

        self.nnunet_raw = self.input_info.pop("nnunet_raw", os.path.join(".", self.work_dir, "nnUNet_raw_data_base"))
        self.nnunet_preprocessed = self.input_info.pop(
            "nnunet_preprocessed", os.path.join(".", self.work_dir, "nnUNet_preprocessed")
        )
        self.nnunet_results = self.input_info.pop(
            "nnunet_results", os.path.join(".", self.work_dir, "nnUNet_trained_models")
        )

        if not os.path.exists(self.nnunet_raw):
            os.makedirs(self.nnunet_raw)

        if not os.path.exists(self.nnunet_preprocessed):

Review comment (suggestion):
Please add a doc-string for `trainer_class_name` after L67. The doc-string, in my opinion, should provide examples such as all the available options, e.g. 1epoch, 5epoch and nnUnetTrainer.

Code after:
        - training for all 20 models (4 configurations by 5 folds) on 2 GPUs

        .. code-block:: bash

            python -m monai.apps.nnunet nnUNetV2Runner train --input_config "./input.yaml" --device_ids "(0,1)"

        - 5-fold training for a single model on 2 GPUs. Here ``configs`` is used to specify the configurations.

        .. code-block:: bash

            python -m monai.apps.nnunet nnUNetV2Runner train --input_config "./input.yaml" \\
                --configs "3d_fullres" \\
                --trainer_class_name "nnUNetTrainer_5epochs" \\
                --export_validation_probabilities True \\
                --device_ids "(0,1)"

        - find the best configuration

        .. code-block:: bash

            python -m monai.apps.nnunet nnUNetV2Runner find_best_configuration --input_config "./input.yaml"

        - predict, ensemble, and post-process

        .. code-block:: bash

            python -m monai.apps.nnunet nnUNetV2Runner predict_ensemble_postprocessing --input_config "./input.yaml"

    """

    def __init__(
        self,
        input_config: Any,
        trainer_class_name: str = "nnUNetTrainer",
        work_dir: str = "work_dir",
        export_validation_probabilities: bool = True,
    ) -> None:
        self.input_info: dict = {}
        self.input_config_or_dict = input_config
        self.trainer_class_name = trainer_class_name
        self.export_validation_probabilities = export_validation_probabilities
        self.work_dir = work_dir

        if isinstance(self.input_config_or_dict, dict):
            self.input_info = self.input_config_or_dict
        elif isinstance(self.input_config_or_dict, str) and os.path.isfile(self.input_config_or_dict):
            self.input_info = ConfigParser.load_config_file(self.input_config_or_dict)
        else:
            raise ValueError(f"{input_config} is not a valid file or dict")

        self.nnunet_raw = self.input_info.pop("nnunet_raw", os.path.join(".", self.work_dir, "nnUNet_raw_data_base"))