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

Code before:
        Args:
            configs: configurations that should be trained.
                Default: ("2d", "3d_fullres", "3d_lowres", "3d_cascade_fullres").
            device_ids: a tuple/list of GPU device IDs to use for the training. Default: None (all available GPUs).
            kwargs: this optional parameter allows you to specify additional arguments defined in the
                ``train_single_model`` method.
        """

        if device_ids is None:
            result = subprocess.run(["nvidia-smi", "--list-gpus"], stdout=subprocess.PIPE)
            output = result.stdout.decode("utf-8")
            num_gpus = len(output.strip().split("\n"))
            device_ids = tuple(range(num_gpus))
        logger.info(f"number of GPUs is {len(device_ids)}, device ids are {device_ids}")
        if len(device_ids) > 1:
            self.train_parallel(
                configs=ensure_tuple(configs),
                device_ids=device_ids,
                trainer_class_name=self.trainer_class_name,
                **kwargs,
            )
        else:
            for cfg in ensure_tuple(configs):
                for _fold in range(self.num_folds):
                    self.train_single_model(
                        config=cfg, fold=_fold, trainer_class_name=self.trainer_class_name, **kwargs
                    )

    def train_parallel_cmd(
        self,
        configs: tuple | str = (M.N_3D_FULLRES, M.N_2D, M.N_3D_LOWRES, M.N_3D_CASCADE_FULLRES),
        device_ids: tuple | None = None,
        **kwargs: Any,
    ) -> list:
        """
        Create the line command for subprocess call for parallel training.

        Args:
            configs: configurations that should be trained.
                Default: ("2d", "3d_fullres", "3d_lowres", "3d_cascade_fullres").
            device_ids: a tuple of GPU device IDs to use for the training. Default: None (all available GPUs).
            kwargs: this optional parameter allows you to specify additional arguments defined in the
                ``train_single_model`` method.
        """
        # unpack compressed files
        folder_names = []
        for root, _, files in os.walk(os.path.join(self.nnunet_preprocessed, self.dataset_name)):
            if any(file.endswith(".npz") for file in files):
                folder_names.append(root)

        from nnunetv2.training.dataloading.utils import unpack_dataset

Review comment (suggestion):
I think the `device_ids` should also be passed in case the user specify a GPU for single-GPU training. Thanks!

Code after:
                fold=fold,
                num_gpus=len(gpu_id),
                trainer_class_name=self.trainer_class_name,
                export_validation_probabilities=self.export_validation_probabilities,
                **kwargs,
            )

    def train(
        self,
        configs: tuple | str = (M.N_3D_FULLRES, M.N_2D, M.N_3D_LOWRES, M.N_3D_CASCADE_FULLRES),
        device_ids: tuple | list | None = None,
        **kwargs: Any,
    ) -> None:
        """
        Run the training for all the models specified by the configurations.
        Note: to set the number of GPUs to use, use ``devices_ids`` instead of the `CUDA_VISIBLE_DEVICES`
        environment variable.

        Args:
            configs: configurations that should be trained.
                Default: ("2d", "3d_fullres", "3d_lowres", "3d_cascade_fullres").
            device_ids: a tuple/list of GPU device IDs to use for the training. Default: None (all available GPUs).
            kwargs: this optional parameter allows you to specify additional arguments defined in the
                ``train_single_model`` method.
        """
        if device_ids is None:
            result = subprocess.run(["nvidia-smi", "--list-gpus"], stdout=subprocess.PIPE)
            output = result.stdout.decode("utf-8")
            num_gpus = len(output.strip().split("\n"))
            device_ids = tuple(range(num_gpus))
        logger.info(f"number of GPUs is {len(device_ids)}, device ids are {device_ids}")
        if len(device_ids) > 1:
            self.train_parallel(configs=ensure_tuple(configs), device_ids=device_ids, **kwargs)
        else:
            for cfg in ensure_tuple(configs):
                for _fold in range(self.num_folds):
                    self.train_single_model(config=cfg, fold=_fold, gpu_id=device_ids, **kwargs)

    def train_parallel_cmd(
        self,
        configs: tuple | str = (M.N_3D_FULLRES, M.N_2D, M.N_3D_LOWRES, M.N_3D_CASCADE_FULLRES),
        device_ids: tuple | list | None = None,
        **kwargs: Any,
    ) -> list:
        """
        Create the line command for subprocess call for parallel training.

        Args:
            configs: configurations that should be trained.
                Default: ("2d", "3d_fullres", "3d_lowres", "3d_cascade_fullres").
            device_ids: a tuple/list of GPU device IDs to use for the training. Default: None (all available GPUs).