o  i@sddlmZddlZddlZddlZddlmZddlmZddl m Z m Z ddl Z ddl mZddlmZddlmZdd lmZdd lmZmZdd lmZdd lmZdd lmZddlmZddl m!Z!m"Z"m#Z#m$Z$ddl%m&Z&m'Z'ee(dZ)e$d\Z*Z+GdddZ,dS)) annotationsN)deepcopy)sleep)Anycast) BundleGen) DataAnalyzer)EnsembleRunner)NNIGen)export_bundle_algo_historyimport_bundle_algo_history) get_logger)algo_to_pickle) ConfigParser) SaveImage)AlgoKeys has_optionlook_up_optionoptional_import) check_kwargs_exist_in_class_initrun_cmd) module_namennic@seZdZUdZded<              dWdXdd Zd!d"Zd#d$ZdYd(d)Z dZd[d-d.Z d\d]d1d2Z d^d_d4d5Z    d`dadd!|jot|_|?i|_@d|_Ad"|j#vrd"|jvr|jd"|j#d"<dSdSdS)#Nz input.yamlz0Input config is not provided, using the default zLoading input config z is not a valid file or dictrz AutoRunner using work directory T)exist_okz cache.yamlr"r$)r"r$r%r&r)r*r-)r r(r+r.r/>modalitydatalistdatarootrzConfig keys are missing r3zDatalist file is not found z!Datalist was copied to work_dir: )datalist_filenamenum_foldzSetting num_fold z based on the input config. based on the input datalist .yamlF)configfilepathfmt sort_keysr4zdatastats.yaml)r6rz HPOGen backend only supports NNIsigmoid)BospathisfilejoinabspathloggerinfodictZ data_src_cfg isinstancerrload_config_file ValueErrorrmakedirsdata_src_cfg_namer r+r-r*cache_filename read_cachecache export_cacher"r$r%r)has_nnir&r(r.r/rr0r'setattr differencekeyslenexistsbasenameshutilcopyfile SameFileErrorinspect_datalist_foldsmax_foldintexport_config_filer4datastats_filenamer5set_training_paramsset_device_infoset_prediction_paramsset_analyze_paramsset_ensemble_method set_num_foldgpu_customizationgpu_customization_specslowerNotImplementedErrorset_hpo_params search_space hpo_tasks)selfrrr r"r$r%r&r(r)r*r+r-r.r/r0param missing_keysr5r6rob/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/apps/auto3dseg/auto_runner.py__init__s&              zAutoRunner.__init__cCsddddd}|jstj|js|St|j}|D] \}}|||q|drCt |dt r;tj|dsCd|d<d|d<|drXt |j dd}t |dkrXd|d<|d rmt |j d d}t |dkrmd|d <|S) af Check if the intermediate result is cached after each step in the current working directory Returns: a dict of cache results. If not_use_cache is set to True, or there is no cache file in the working directory, the result will be ``empty_cache`` in which all ``has_cache`` keys are set to False. FN)r" datastatsr$r%r"rrr$Z only_trainedrr%T)r*r?r@rArLrrHitems setdefaultrGrr rrT)rl empty_cacherNkvhistoryZtrained_historyrororprM\s&    zAutoRunner.read_cachecKs(|j|tj|j|jdddddS)zQ Save the cache state as ``cache.yaml`` in the working directory r9NF)r<default_flow_styler=)rNupdaterr]rL)rlr0rororprOs  zAutoRunner.export_cacher5returnr\c Cst|}d|vrtdt|dd|dD}t|dkrEt|d}td|d|d tt||krCtd |d|Sd |vrt|d dkrtd |dD]}d|d <qZ|d D]}d|d <qedd|d D}tdt|d||dD]}|d|vrd|d <||d=q|dt | |d<tj ||dddd}|Sd}t d|d|dddlm}||ddd}t||dD]\} \} } | D] } | |d| d <qqtj ||ddd|S)a Returns number of folds in the datalist file, and assigns fold numbers if not provided. Args: datalist_filename: path to the datalist file. Notes: If the fold key is not provided, it auto generates 5 folds assignments in the training key list. If validation key list is available, then it assumes a single fold validation. trainingz#Datalist files has no training key:cSs g|] }d|vrt|dqS)fold)r\.0drororp  z5AutoRunner.inspect_datalist_folds..rzFound num_fold r7r8z*Fold numbers are not continuous from 0 to validationzUNo fold numbers provided, attempting to use a single fold based on the validation keyr~cSsi|] }d|vr|d|qS)labelrorrororp sz5AutoRunner.inspect_datalist_folds..zFound z8 items in the validation key, saving updated datalist torjson)r<indentz Datalist has no folds specified z...Generating z[ folds randomly.Please consider presaving fold numbers beforehand for repeated experiments.)KFoldT)Zn_splitsshuffle random_state)rrHrIrrTmaxrDrEsetlistvaluesr]warningswarnZsklearn.model_selectionr enumeratesplit) rlr5r3Z fold_listr6rZ val_labelsrkfi_Z valid_idxvirororprZsV   +         z!AutoRunner.inspect_datalist_foldsrerfdict[str, Any] | NonecCs||_|dur ||_|S)a Set options for GPU-based parameter customization/optimization. Args: gpu_customization: the switch to determine automatically customize/optimize bundle script/config parameters for each bundleAlgo based on gpus. Custom parameters are obtained through dummy training to simulate the actual model training process and hyperparameter optimization (HPO) experiments. gpu_customization_specs (optional): the dictionary to enable users overwrite the HPO settings. user can overwrite part of variables as follows or all of them. The structure is as follows. .. code-block:: python gpu_customization_specs = { 'ALGO': { 'num_trials': 6, 'range_num_images_per_batch': [1, 20], 'range_num_sw_batch_size': [1, 20] } } ALGO: the name of algorithm. It could be one of algorithm names (e.g., 'dints') or 'universal' which would apply changes to all algorithms. Possible options are - {``"universal"``, ``"dints"``, ``"segresnet"``, ``"segresnet2d"``, ``"swinunetr"``}. num_trials: the number of HPO trials/experiments to run. range_num_images_per_batch: the range of number of images per mini-batch. range_num_sw_batch_size: the range of batch size in sliding-window inferer. N)rerf)rlrerfrororpset_gpu_customizations!z AutoRunner.set_gpu_customizationrr6cCsD|dkr td|||jkrtd|jd|jd||_|S)z Set the number of cross validation folds for all algos. Args: num_fold: a positive integer to define the number of folds. rzEnum_fold is expected to be an integer greater than zero. Now it gets z1num_fold is greater than the maximum fold number z in r8)rIr[r5r6)rlr6rororprds zAutoRunner.set_num_foldparamscCs0|durt|ni|_d|jvrtdt|S)a Set the training params for all algos. Args: params: a dict that defines the overriding key-value pairs during training. The overriding method is defined by the algo class. Examples: For BundleAlgo objects, the training parameter to shorten the training time to a few epochs can be {"num_epochs": 2, "num_epochs_per_validation": 1} NCUDA_VISIBLE_DEVICESz]CUDA_VISIBLE_DEVICES is deprecated from 'set_training_params'. Use 'set_device_info' instead.)r train_paramsrrDeprecationWarningrlrrororpr_s zAutoRunner.set_training_paramscuda_visible_deviceslist[int] | str | None num_nodes int | Nonemn_start_method cmd_prefixcCsNi|_|dur tjd}|dur,dddttjD|jd<tj|jd<n:t |t rA||jd<t | d|jd<n%t |t tfr]ddd|D|jd<t ||jd<n td|d |dursttjd d }||jd <|durtjd d }||jd <|durtjdd}||jd<|durtd|d|S)a Set the device related info Args: cuda_visible_devices: define GPU ids for data analyzer, training, and ensembling. List of GPU ids [0,1,2,3] or a string "0,1,2,3". Default using env "CUDA_VISIBLE_DEVICES" or all devices available. num_nodes: number of nodes for training and ensembling. Default using env "NUM_NODES" or 1 if "NUM_NODES" is unset. mn_start_method: multi-node start method. Autorunner will use the method to start multi-node processes. Default using env "MN_START_METHOD" or 'bcprun' if "MN_START_METHOD" is unset. cmd_prefix: command line prefix for subprocess running in BundleAlgo and EnsembleRunner. Default using env "CMD_PREFIX" or None, examples are: - single GPU/CPU or multinode bcprun: "python " or "/opt/conda/bin/python3.9 ", - single node multi-GPU running "torchrun --nnodes=1 --nproc_per_node=2 " If user define this prefix, please make sure --nproc_per_node matches cuda_visible_device or os.env['CUDA_VISIBLE_DEVICES']. Also always set --nnodes=1. Set num_nodes for multi-node. Nr,cSg|]}t|qSrorrxrororpr=z.AutoRunner.set_device_info.. n_devicescSrrorrrororprCrz%Wrong format of cuda_visible_devices z, devices not setZ NUM_NODESrZMN_START_METHODZbcprunZ CMD_PREFIXz*Using user defined command running prefix z, will override other settings)device_settingr?environgetrBrangetorchcuda device_countrGrrTrrtuplerDwarningr\rE)rlrrrrrororpr`s2 $     zAutoRunner.set_device_infoAlgoEnsembleBestByFoldensemble_method_namecKs"t|ddgd|_|j||S)a Set the bundle ensemble method name and parameters for save image transform parameters. Args: ensemble_method_name: the name of the ensemble method. Only two methods are supported "AlgoEnsembleBestN" and "AlgoEnsembleBestByFold". kwargs: the keyword arguments used to define the ensemble method. Currently only ``n_best`` for ``AlgoEnsembleBestN`` is supported. AlgoEnsembleBestNr) supported)rrr0r{)rlrr0rororprcYs   zAutoRunner.set_ensemble_methodcKs0tt|\}}|r|j||St|d)a# Set the ensemble output transform. Args: kwargs: image writing parameters for the ensemble inference. The kwargs format follows SaveImage transform. For more information, check https://docs.monai.io/en/stable/transforms.html#saveimage. z are not supported in monai.transforms.SaveImage,Check https://docs.monai.io/en/stable/transforms.html#saveimage for more information.)rrr0r{rI)rlr0Zare_all_args_present extra_argsrororpset_image_save_transformjs  z#AutoRunner.set_image_save_transformcCs |dur t||_|Si|_|S)a Set the prediction params for all algos. Args: params: a dict that defines the overriding key-value pairs during prediction. The overriding method is defined by the algo class. Examples: For BundleAlgo objects, this set of param will specify the algo ensemble to only inference the first two files in the testing datalist {"file_slices": slice(0, 2)} N)r pred_paramsrrororprasz AutoRunner.set_prediction_paramscCs&|dur ddd|_|St||_|S)z Set the data analysis extra params. Args: params: a dict that defines the overriding key-value pairs during training. The overriding method is defined by the algo class. NFr)Zdo_ccpdevice)rrrrororprbs   zAutoRunner.set_analyze_paramscCs|dur |j|_|S||_|S)a Set parameters for the HPO module and the algos before the training. It will attempt to (1) override bundle templates with the key-value pairs in ``params`` (2) change the config of the HPO module (e.g. NNI) if the key is found to be one of: - "trialCodeDirectory" - "trialGpuNumber" - "trialConcurrency" - "maxTrialNumber" - "maxExperimentDuration" - "tuner" - "trainingService" and (3) enable the dry-run mode if the user would generate the NNI configs without starting the NNI service. Args: params: a dict that defines the overriding key-value pairs during instantiation of the algo. For BundleAlgo, it will override the template config filling. Notes: Users can set ``nni_dry_run`` to ``True`` in the ``params`` to enable the dry-run mode for the NNI backend. N)r hpo_paramsrrororpriszAutoRunner.set_hpo_paramsrjdict[str, Any]cCsXd}|D]\}}d|vrt|d|d|d|t|d9}q||_||_|S)a$ Set the search space for NNI parameter search. Args: search_space: hyper parameter search space in the form of dict. For more information, please check NNI documentation: https://nni.readthedocs.io/en/v2.2/Tutorial/SearchSpaceSpec.html . r_valuez key z value z has not _value)rtrIrTrjrk)rlrjZvalue_combinationsrwrxrororpset_nni_search_spaceszAutoRunner.set_nni_search_spacerylist[dict[str, Any]]NonecCsn|D]2}|tj}t|jdr||j|jn||j|}ttj|i}t |fd|j i|qdS)a Train the Algos in a sequential scheme. The order of training is randomized. Args: history: the history of generated Algos. It is a list of dicts. Each element has the task name (e.g. "dints_0" for dints network in fold 0) as the key and the algo object as the value. After the training, the algo object with the ``best_metric`` will be saved as a pickle file. Note: The final results of the model training will be written to all the generated algorithm's output folders under the working directory. The results include the model checkpoints, a progress.yaml, accuracies in CSV and a pickle file of the Algo object. r template_pathN) rALGOrr%rr get_scorerSCORErr)rlry algo_dictalgoaccZalgo_meta_datarororp_train_algo_in_sequences   z"AutoRunner._train_algo_in_sequencec Csdtjdddddiddd d }tt|jdd }|jd d }|D]}|tj }|tj }t ||jd}| } t |} |jD] } | | vrO|j| | | <qB| d|i| d|jid| d|j} | d| itjtj|j|d} tj| | dddt|jtt|d}d| d}|rtd|q$t|ddtt|jdd }|||krtdtt|jdd }|||ksd}t|ddtd||}q$dS)a! Train the Algos using HPO. Args: history: the history of generated Algos. It is a list of dicts. Each element has the task name (e.g. "dints_0" for dints network in fold 0) as the key and the algo object as the value. After the training, the algo object with the ``best_metric`` will be saved as a pickle file. Note: The final results of the model training will not be written to all the previously generated algorithm's output folders. Instead, HPO will generate a new algo during the searching, and the new algo will be saved under the working directory with a different format of the name. For example, if the searching space has "learning_rate", the result of HPO will be written to a folder name with original task name and the param (e.g. "dints_0_learning_rate_0.001"). The results include the model checkpoints, a progress.yaml, accuracies in CSV and a pickle file of the Algo object. r8r 1hnameZ GridSearchlocalT)platformZ useActiveGpu)ZtrialCodeDirectoryZtrialGpuNumberZtrialConcurrencymaxTrialNumberZmaxExperimentDurationZtunerZtrainingServicersZ nni_dry_runF)rrZexperimentNamerjz/python -m monai.apps.auto3dseg NNIGen run_algo  Z trialCommandz_nni_config.yamlr9N)r<rzrznnictl create --config z --port 8088z;AutoRunner HPO is in dry-run mode. Please manually launch: )checkznnictl stop --allzNNI completes HPO on )rrrrTr rrpoprIDrr Zget_obj_filenamerr{rjr?r@rCrBrr]minrkrr\rDrErrr)rlryZdefault_nni_configZlast_total_tasksZ mode_dry_runrrrZnni_genZ obj_filenameZ nni_configkeyZ trial_cmdZnni_config_filenameZ max_trialcmdZ n_trainingsrororp_train_algo_in_nnisT       zAutoRunner._train_algo_in_nnic Cs2|jr0|jdur0tdt|j|jfd|ji|j}|d}t j |j d|jdntd|j rtj|jsHtd|jdt|j|j|j|j|j|j|jd }|jrm|j|j|j|j|j|jd n |j|j|j|jd |}t||j dd ntd |j du}|j s|r|j!dst"|jdd}t#|dkrtd|jd|rdd|D}|rtd|ddd|D}t#|dkr|j$s|%|n|&||j ddntd|j'rt(d|j|j|j|j)t*|j+ddkd|j,|j-}|.|j+tddS) z- Run the AutoRunner pipeline NzRunning data analysis... output_pathT)r"rrzSkipping data analysis...z"Could not find the datastats file z=. Possibly the required data analysis step was not completed.)r Z algo_pathr+Zdata_stats_filenamerKr.r/)r6rerfr-)r6r-)r$z Skipping algorithm generation...r%Frsrz#Could not find training scripts in zE. Possibly the required algorithms generation step was not completed.cSs g|] }|tjr|tjqSro)r IS_TRAINEDrrhrororprmrz"AutoRunner.run..zSkipping already trained algos z3.Set option train=True to always retrain all algos.cSsg|] }|tjs|qSro)rrrrororprss)r%zSkipping algorithm training...rr)rKrr6rZmgpuz-Auto3Dseg pipeline is completed successfully.ro)/r"rrDrErr5r4r^Zget_all_case_statsrrrvrOr$r?r@rArIrr rr+rKr.r/regenerater6rfr-Z get_historyr r%rNr rTr&rrr)r rr\rr0rrun)rldaZbundle_generatorryZauto_train_choiceZ skip_algosZensemble_runnerrororpr,s               zAutoRunner.run)rNNNNNFrTFNTNN)rrrrr r!r"r#r$r#r%r#r&r'r(rr)r'r*r'r+r,r-r'r.r,r/r,r0r)r5rr|r\)FN)rer'rfrr|r)r)r6r\r|r)N)rrr|r)NNNN) rrrrrr,rr,r|r)r)rrr0rr|r)r0rr|r)rjrr|r)ryrr|r)__name__ __module__ __qualname____doc____annotations__rqrMrOrZrrdr_r`rcrrarbrirrrrrorororpr(sR (  % E '  ;        Cr)- __future__rr?rWrcopyrtimertypingrrrZmonai.apps.auto3dseg.bundle_genrZ"monai.apps.auto3dseg.data_analyzerrZ%monai.apps.auto3dseg.ensemble_builderr Zmonai.apps.auto3dseg.hpo_genr Zmonai.apps.auto3dseg.utilsr r monai.apps.utilsr Zmonai.auto3dseg.utilsrZ monai.bundlermonai.transformsr monai.utilsrrrrmonai.utils.miscrrrrDrrPrrorororps,