U
    {Ph                     @  s   d dl mZ d dlZd dlmZ d dlmZmZ d dlm	Z	m
Z
 d dlmZ d dlmZmZmZ edejed	\ZZerd d
lmZ nedejeddd\ZZG dd dZdS )    )annotationsN)Callable)TYPE_CHECKINGAny)ReduceLROnPlateau_LRScheduler)
IgniteInfo)ensure_tuplemin_versionoptional_importzignite.engineEvents)Enginer   	decorator)as_typec                   @  sV   e Zd ZdZddddd fddddd	d
dddZdd
dddZdd
dddZdS )LrScheduleHandlerzS
    Ignite handler to update the Learning Rate based on PyTorch LR scheduler.
    TNc                 C  s   dS )N r   )enginer   r   W/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/handlers/lr_schedule_handler.py<lambda>+       zLrScheduleHandler.<lambda>z _LRScheduler | ReduceLROnPlateauboolz
str | NonezCallable[[Engine], Any]None)lr_schedulerprint_lrnameepoch_levelstep_transformreturnc                 C  sL   || _ || _t|| _|| _t|s<tdt|j	 d|| _
|| _dS )a  
        Args:
            lr_scheduler: typically, lr_scheduler should be PyTorch
                lr_scheduler object. If customized version, must have `step` and `get_last_lr` methods.
            print_lr: whether to print out the latest learning rate with logging.
            name: identifier of logging.logger to use, if None, defaulting to ``engine.logger``.
            epoch_level: execute lr_scheduler.step() after every epoch or every iteration.
                `True` is epoch level, `False` is iteration level.
            step_transform: a callable that is used to transform the information from `engine`
                to expected input data of lr_scheduler.step() function if necessary.

        Raises:
            TypeError: When ``step_transform`` is not ``callable``.

        z'step_transform must be callable but is .N)r   r   logging	getLoggerloggerr   callable	TypeErrortype__name__r   _name)selfr   r   r   r   r   r   r   r   __init__%   s    zLrScheduleHandler.__init__r   )r   r   c                 C  s:   | j dkr|j| _| jr(|tj|  n|tj|  dS )g
        Args:
            engine: Ignite Engine, it can be a trainer, validator or evaluator.
        N)r&   r!   r   add_event_handlerr   EPOCH_COMPLETEDITERATION_COMPLETED)r'   r   r   r   r   attachF   s
    
zLrScheduleHandler.attachc                 C  s>   t | |}| jj|  | jr:| jd| jjd   dS )r)   zCurrent learning rate: r   N)r	   r   r   stepr   r!   info_last_lr)r'   r   argsr   r   r   __call__R   s    zLrScheduleHandler.__call__)r%   
__module____qualname____doc__r(   r-   r2   r   r   r   r   r       s   !r   )
__future__r   r   collections.abcr   typingr   r   torch.optim.lr_schedulerr   r   monai.configr   monai.utilsr	   r
   r   OPT_IMPORT_VERSIONr   _ignite.enginer   r   r   r   r   r   <module>   s"       
