U
    Ph"A                     @  s  d dl mZ d dlZd dlmZ d dlmZ d dlmZm	Z	 d dl
mZ d dlmZ d dlmZ d d	lmZmZmZmZ d d
lmZ d dlmZ d dlmZ d dlmZ ed\ZZed\ZZ ee!dZ"dddgZ#G dd deZ$G dd de$Z%G dd de$Z&dS )    )annotationsN)abstractmethod)deepcopy)Anycast)warn)
BundleAlgo)
get_logger)AlgoAlgoGenalgo_from_picklealgo_to_pickle)ConfigParser)PathLike)optional_import)AlgoKeysnnioptuna)module_nameHPOGenNNIGen	OptunaGenc                   @  s@   e Zd ZdZedd Zedd Zedd Zedd	 Zd
S )r   a>  
    The base class for hyperparameter optimization (HPO) interfaces to generate algos in the Auto3Dseg pipeline.
    The auto-generated algos are saved at their ``output_path`` on the disk. The files in the ``output_path``
    may contain scripts that define the algo, configuration files, and pickle files that save the internal states
    of the algo before/after the training. Compared to the BundleGen class, HPOGen generates Algo on-the-fly, so
    training and algo generation may be executed alternatively and take a long time to finish the generation process.

    c                 C  s   t dS )z Get the hyperparameter from HPO.NNotImplementedErrorself r   Q/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/apps/auto3dseg/hpo_gen.pyget_hyperparameters-   s    zHPOGen.get_hyperparametersc                 O  s   t dS )zHUpdate Algo parameters according to the hyperparameters to be evaluated.Nr   r   argskwargsr   r   r   update_params2   s    zHPOGen.update_paramsc                 C  s   t dS )z$Report the evaluated results to HPO.Nr   r   r   r   r   	set_score7   s    zHPOGen.set_scorec                 O  s   t dS )zDInterface for launch the training given the fetched hyperparameters.Nr   r   r   r   r   run_algo<   s    zHPOGen.run_algoN)	__name__
__module____qualname____doc__r   r   r"   r#   r$   r   r   r   r   r   #   s   	


c                   @  s   e Zd ZdZd dddddZdd	 Zd
d Zdd ZdddddZdd Z	d!dddddZ
dd Zd"dddddddZdS )#r   u  
    Generate algorithms for the NNI to automate hyperparameter tuning. The module has two major interfaces:
    ``__init__`` which prints out how to set up the NNI, and a trialCommand function ``run_algo`` for the NNI library to
    start the trial of the algo. More about trialCommand function can be found in ``trail code`` section in NNI webpage
    https://nni.readthedocs.io/en/latest/tutorials/hpo_quickstart_pytorch/main.html .

    Args:
        algo: an Algo object (e.g. BundleAlgo) with defined methods: ``get_output_path`` and train
            and supports saving to and loading from pickle files via ``algo_from_pickle`` and ``algo_to_pickle``.
        params: a set of parameter to override the algo if override is supported by Algo subclass.

    Examples::

        The experiment will keep generating new folders to save the model checkpoints, scripts, and configs if available.
        ├── algorithm_templates
        │   └── unet
        ├── unet_0
        │   ├── algo_object.pkl
        │   ├── configs
        │   └── scripts
        ├── unet_0_learning_rate_0.01
        │   ├── algo_object.pkl
        │   ├── configs
        │   ├── model_fold0
        │   └── scripts
        └── unet_0_learning_rate_0.1
            ├── algo_object.pkl
            ├── configs
            ├── model_fold0
            └── scripts

        .. code-block:: python
            # Bundle Algorithms are already generated by BundleGen in work_dir
            import_bundle_algo_history(work_dir, only_trained=False)
            algo_dict = self.history[0]  # pick the first algorithm
            algo_name = algo_dict[AlgoKeys.ID]
            onealgo = algo_dict[AlgoKeys.ALGO]
            nni_gen = NNIGen(algo=onealgo)
            nni_gen.print_bundle_algo_instruction()

    Notes:
        The NNIGen will prepare the algorithms in a folder and suggest a command to replace trialCommand in the experiment
        config. However, NNIGen will not trigger NNI. User needs to write their NNI experiment configs, and then run the
        NNI command manually.
    NAlgo | Nonedict | None)algoparamsc                 C  s   |  d| _ d| _|d k	rt|tr|d kr2|| _qt|| _tj|	 d }tj
|	 }|ddi | jj||f| n|| _t| j| jjd| _d S N Z	_overridefill_with_datastatsFtemplate_path)hintobj_filename
isinstancer   r+   r   ospathbasenameget_output_pathdirnameupdateexport_to_diskr   r1   r   r+   r,   nameoutput_folderr   r   r   __init__q   s    

zNNIGen.__init__c                 C  s   | j S )z5Return the filename of the dumped pickle algo object.r3   r   r   r   r   get_obj_filename   s    zNNIGen.get_obj_filenamec                 C  s   d}t d t d t d t | d| j d t d t d t d	tt| jj d
 t d| j  d t d t | d t d dS )zH
        Print how to write the trial commands for Bundle Algo.
        z/python -m monai.apps.auto3dseg NNIGen run_algo z============================================================================================================================================z#If NNI will run in your local env: zB1. Add the following line to the trialCommand in your NNI config:  z {result_dir}z--------------------------------------------------------------------------------------------------------------------------------------------z!If NNI will run in a remote env: z'1. Copy the algorithm_templates folder z+ to remote {remote_algorithm_templates_dir}z2. Copy the older z( to the remote machine {remote_algo_dir}zDThen add the following line to the trialCommand in your NNI config: z@ {remote_algo_dir} {result_dir} {remote_algorithm_templates_dir}N)loggerinfor3   r   r   r+   r1   r8   )r   r2   r   r   r   print_bundle_algo_instruction   s    





z$NNIGen.print_bundle_algo_instructionc                 C  s   t rt S td i S )zK
        Get parameter for next round of training from NNI server.
        ?NNI is not detected. The code will continue to run without NNI.)has_nnir   Zget_next_parameterr   r   r   r   r   r      s    zNNIGen.get_hyperparametersdictNoner,   returnc                 C  s
   || _ dS )z
        Translate the parameter from monai bundle to meet NNI requirements.

        Args:
            params: a dict of parameters.
        Nr,   r   r,   r   r   r   r"      s    zNNIGen.update_paramsc                 C  s   d dd | j D pdS )
        Get the identifier of the current experiment. In the format of listing the searching parameter name and values
        connected by underscore in the file name.
        r.   c                 s  s"   | ]\}}d | d | V  qdS _Nr   .0kvr   r   r   	<genexpr>   s     z%NNIGen.get_task_id.<locals>.<genexpr>_Nonejoinr,   itemsr   r   r   r   get_task_id   s    zNNIGen.get_task_id.strr>   rK   c                 C  s~   |   }tj| j }tj||| }tj|d| _t| jt	rb| jj
||| |dd nt| j| t| dS )
        Generate the record for each Algo. If it is a BundleAlgo, it will generate the config files.

        Args:
            output_folder: the directory nni will save the results to.
        algo_object.pklF)bundle_rootr/   NrZ   r5   r6   r7   r+   r8   rX   r3   r4   r   r;   r   export_config_filer,   rC   rD   r   r>   Ztask_idZtask_prefix
write_pathr   r   r   generate   s       zNNIGen.generatec                 C  s   t rt| ntd dS )z/
        Report the acc to NNI server.
        rF   N)rG   r   Zreport_final_resultr   r   accr   r   r   r#      s    zNNIGen.set_scorePathLike | Noner3   r>   r1   rK   c                 C  s   t j|st| dt||d\| _}|  }| | | | | j	| j
 | j }ttj|i}t| jfd| jji| | | dS am  
        The python interface for NNI to run.

        Args:
            obj_filename: the pickle-exported Algo object.
            output_folder: the root path of the algorithms templates.
            template_path: the algorithm_template. It must contain algo.py in the follow path:
                ``{algorithm_templates_dir}/{network}/scripts/algo.py``
        z is not foundr0   r1   Nr5   r6   isfile
ValueErrorr   r+   r   r"   re   trainr,   	get_scorer\   r   SCOREr   r1   r#   r   r3   r>   r1   algo_meta_datar,   rg   r   r   r   r$      s    



zNNIGen.run_algo)NN)r[   )r[   N)r%   r&   r'   r(   r?   rA   rE   r   r"   rZ   re   r#   r$   r   r   r   r   r   B   s   .			c                   @  s   e Zd ZdZd$ddddddZd	d
 Zdd Zdd Zdd Zd%ddddddddZ	dddddZ
dd Zd&ddddd Zd'ddddd!d"d#ZdS )(r   u  
    Generate algorithms for the Optuna to automate hyperparameter tuning. Please refer to NNI and Optuna
    (https://optuna.readthedocs.io/en/stable/) for more information. Optuna has different running scheme
    compared to NNI. The hyperparameter samples come from a trial object (trial.suggest...) created by Optuna,
    so OptunaGen needs to accept this trial object as input. Meanwhile, Optuna calls OptunaGen,
    thus OptunaGen.__call__() should return the accuracy. Use functools.partial to wrap OptunaGen
    for addition input arguments.

    Args:
        algo: an Algo object (e.g. BundleAlgo). The object must at least define two methods: get_output_path and train
            and supports saving to and loading from pickle files via ``algo_from_pickle`` and ``algo_to_pickle``.
        params: a set of parameter to override the algo if override is supported by Algo subclass.

    Examples::

        The experiment will keep generating new folders to save the model checkpoints, scripts, and configs if available.
        ├── algorithm_templates
        │   └── unet
        ├── unet_0
        │   ├── algo_object.pkl
        │   ├── configs
        │   └── scripts
        ├── unet_0_learning_rate_0.01
        │   ├── algo_object.pkl
        │   ├── configs
        │   ├── model_fold0
        │   └── scripts
        └── unet_0_learning_rate_0.1
            ├── algo_object.pkl
            ├── configs
            ├── model_fold0
            └── scripts

    Notes:
        Different from NNI and NNIGen, OptunaGen and Optuna can be ran within the Python process.

    Nr)   r*   rI   )r+   r,   rK   c                 C  s   |  d| _ |d k	rt|tr||d kr,|| _qt|| _tj| d }tj	| }|
ddi | jj||f| n|| _t| j| jjd| _ d S r-   )r3   r4   r   r+   r   r5   r6   r7   r8   r9   r:   r;   r   r1   r<   r   r   r   r?     s    

zOptunaGen.__init__c                 C  s   | j S )z(Return the dumped pickle object of algo.r@   r   r   r   r   rA   -  s    zOptunaGen.get_obj_filenamec                 C  s2   t r"td d| jdddiS td i S dS )z
        Get parameter for next round of training from optuna trial object.
        This function requires user rewrite during usage for different search space.
        z2Please rewrite this code by creating a child classZlearning_rateg-C6?g?zEOptuna is not detected. The code will continue to run without Optuna.N)
has_optunarC   rD   trialZsuggest_floatr   r   r   r   r   r   1  s
    
zOptunaGen.get_hyperparametersc                 C  s
   || _ dS )zSet the accuracy scoreN)rg   rf   r   r   r   r#   =  s    zOptunaGen.set_scorec                 C  s
   || _ dS )zSet the Optuna trialN)rt   )r   rt   r   r   r   	set_trialA  s    zOptunaGen.set_trialr[   r   r\   rh   )rt   r3   r>   r1   rK   c                 C  s   |  | | ||| | jS )a  
        Callable that Optuna will use to optimize the hyper-parameters

        Args:
            obj_filename: the pickle-exported Algo object.
            output_folder: the root path of the algorithms templates.
            template_path: the algorithm_template. It must contain algo.py in the follow path:
                ``{algorithm_templates_dir}/{network}/scripts/algo.py``
        )ru   r$   rg   )r   rt   r3   r>   r1   r   r   r   __call__E  s    
zOptunaGen.__call__rH   rJ   c                 C  s
   || _ dS )zu
        Translate the parameter from monai bundle.

        Args:
            params: a dict of parameters.
        NrL   rM   r   r   r   r"   U  s    zOptunaGen.update_paramsc                 C  s   d dd | j D pdS )rN   r.   c                 s  s"   | ]\}}d | d | V  qdS rO   r   rQ   r   r   r   rU   c  s     z(OptunaGen.get_task_id.<locals>.<genexpr>rV   rW   r   r   r   r   rZ   ^  s    zOptunaGen.get_task_idr]   c                 C  s|   |   }tj| j }tj||| }tj|d| _t| jt	r`| jj
||| dd nt| j| t| dS )r^   r_   F)r/   Nra   rc   r   r   r   re   e  s    zOptunaGen.generateri   c                 C  s   t j|st| dt||d\| _}|  }| | | | | j	| j
 | j }ttj|i}t| jfd| jji| | | dS rj   rk   rq   r   r   r   r$   w  s    



zOptunaGen.run_algo)NN)r[   N)r[   )r[   N)r%   r&   r'   r(   r?   rA   r   r#   ru   rv   r"   rZ   re   r$   r   r   r   r   r      s   &   	)'
__future__r   r5   abcr   copyr   typingr   r   warningsr   monai.apps.auto3dseg.bundle_genr   monai.apps.utilsr	   monai.auto3dsegr
   r   r   r   monai.bundle.config_parserr   monai.configr   monai.utilsr   monai.utils.enumsr   r   rG   r   rs   r%   rC   __all__r   r   r   r   r   r   r   <module>   s(   

 1