o / iU@sHddlmZddlZddlZddlZddlmZddlmZm Z m Z ddl Z ddl Z ddlmZddlmZddlmZddlmZddlmZdd lmZmZdd lmZmZmZerkddlm Z!d Z"ddl#Z#d Z$n ed \Z!Z"ed \Z#Z$dgZ%GdddZ&Gddde&Z'Gddde&Z(dddZ)dddZ*GdddZ+dS)) annotationsN)partial) TYPE_CHECKINGAnyCallable) Optimizer)DEFAULT_PROTOCOL) DataLoader) eval_mode) ExponentialLRLinearLR) StateCachercopy_to_deviceoptional_importTzmatplotlib.pyplottqdmLearningRateFinderc@s:eZdZddd Zed d Zd d ZddZddZdS)DataLoaderIter data_loaderr image_extractorrlabel_extractorreturnNonecCs>t|tstdt|d||_t||_||_||_dS)NzLoader has unsupported type: z1. Expected type was `torch.utils.data.DataLoader`) isinstancer ValueErrortyperiter _iteratorrrselfrrrr\/home/dell461/cl/sdc2/last_ska_mid/HISourceFinder-master-l/src/monai/optimizers/lr_finder.py__init__/s   zDataLoaderIter.__init__cCs|jjSN)rdatasetrrrr r#9szDataLoaderIter.datasetcCs||}||}||fSr")rr)r batch_dataimageslabelsrrr inputs_labels_from_batch=s  z'DataLoaderIter.inputs_labels_from_batchcCs|Sr"rr$rrr __iter__BszDataLoaderIter.__iter__cCst|j}||Sr")nextrr()rbatchrrr __next__Es  zDataLoaderIter.__next__Nrr rrrrrr) __name__ __module__ __qualname__r!propertyr#r(r)r,rrrr r-s   rcs*eZdZ ddfd d Zd dZZS)TrainDataLoaderIterTrr rrr auto_resetboolrrcst|||||_dSr")superr!r3)rrrrr3 __class__rr r!Ls zTrainDataLoaderIter.__init__cCshzt|j}||\}}W||fSty3|jst|j|_t|j}||\}}Y||fSwr")r*rr( StopIterationr3rr)rr+inputsr'rrr r,Rs    zTrainDataLoaderIter.__next__T) rr rrrrr3r4rr)r.r/r0r!r, __classcell__rrr6r r2Jsr2cs6eZdZdZdfd d Zd d Zfd dZZS)ValDataLoaderIteraThis iterator will reset itself **only** when it is acquired by the syntax of normal `iterator`. That is, this iterator just works like a `torch.data.DataLoader`. If you want to restart it, you should use it like: ``` loader_iter = ValDataLoaderIter(data_loader) for batch in loader_iter: ... # `loader_iter` should run out of values now, you can restart it by: # 1. the way we use a `torch.data.DataLoader` for batch in loader_iter: # __iter__ is called implicitly ... # 2. passing it into `iter()` manually loader_iter = iter(loader_iter) # __iter__ is called by `iter()` ``` rr rrrrrcs&t|||t|j|_d|_dSNr)r5r!lenr run_limit run_counterrr6rr r!us  zValDataLoaderIter.__init__cCs"|j|jkrt|j|_d|_|Sr=)r@r?rrrr$rrr r)zs  zValDataLoaderIter.__iter__cs|jd7_tS)N)r@r5r,r$r6rr r,s zValDataLoaderIter.__next__r-)r.r/r0__doc__r!r)r,r;rrr6r r<`s r<xrr torch.TensorcC"t|tr |d}|S|d}|S)z3Default callable for getting image from batch data.imagerrdictrCoutrrr default_image_extractorrKcCrE)z3Default callable for getting label from batch data.labelrArGrIrrr default_label_extractorrLrNc @seZdZdZddddeedfdUddZdVddZdee ddddd d!d"ddf dWd6d7Z dXd:d;Z d>> lr_finder = LearningRateFinder(net, optimizer, criterion) >>> lr_finder.range_test(data_loader, end_lr=100, num_iter=100) >>> lr_finder.get_steepest_gradient() >>> lr_finder.plot() # to inspect the loss-learning rate graph Example (Leslie Smith's approach): >>> lr_finder = LearningRateFinder(net, optimizer, criterion) >>> lr_finder.range_test(train_loader, val_loader=val_loader, end_lr=1, num_iter=100, step_mode="linear") Gradient accumulation is supported; example: >>> train_data = ... # prepared dataset >>> desired_bs, real_bs = 32, 4 # batch size >>> accumulation_steps = desired_bs // real_bs # required steps for accumulation >>> data_loader = torch.utils.data.DataLoader(train_data, batch_size=real_bs, shuffle=True) >>> acc_lr_finder = LearningRateFinder(net, optimizer, criterion) >>> acc_lr_finder.range_test(data_loader, end_lr=10, num_iter=100, accumulation_steps=accumulation_steps) By default, image will be extracted from data loader with x["image"] and x[0], depending on whether batch data is a dictionary or not (and similar behaviour for extracting the label). If your data loader returns something other than this, pass a callable function to extract it, e.g.: >>> image_extractor = lambda x: x["input"] >>> label_extractor = lambda x: x[100] >>> lr_finder = LearningRateFinder(net, optimizer, criterion) >>> lr_finder.range_test(train_loader, val_loader, image_extractor, label_extractor) References: Modified from: https://github.com/davidtvs/pytorch-lr-finder. Cyclical Learning Rates for Training Neural Networks: https://arxiv.org/abs/1506.01186 NTFmodel nn.Module optimizerr criteriontorch.nn.Moduledevicestr | torch.device | None memory_cacher4 cache_dir str | Noneamp pickle_moduletypes.ModuleTypepickle_protocolintverboserrc Cs||_|||_||_ggd|_||_||_||_| |_t |j j |_ t |||| d|_|jd|j|jd|j|rL||_ dS|j |_ dS)a8Constructor. Args: model: wrapped model. optimizer: wrapped optimizer. criterion: wrapped loss function. device: device on which to test. run a string ("cpu" or "cuda") with an optional ordinal for the device type (e.g. "cuda:X", where is the ordinal). Alternatively, can be an object representing the device on which the computation will take place. Default: None, uses the same device as `model`. memory_cache: if this flag is set to True, `state_dict` of model and optimizer will be cached in memory. Otherwise, they will be saved to files under the `cache_dir`. cache_dir: path for storing temporary files. If no path is specified, system-wide temporary directory is used. Notice that this parameter will be ignored if `memory_cache` is True. amp: use Automatic Mixed Precision pickle_module: module used for pickling metadata and objects, default to `pickle`. this arg is used by `torch.save`, for more details, please check: https://pytorch.org/docs/stable/generated/torch.save.html#torch.save. pickle_protocol: can be specified to override the default protocol, default to `2`. this arg is used by `torch.save`, for more details, please check: https://pytorch.org/docs/stable/generated/torch.save.html#torch.save. verbose: verbose output Returns: None lrloss) in_memoryrWrZr\rOrQN)rQ_check_for_schedulerrOrRhistoryrVrWrYr^r* parametersrT model_devicer state_cacherstore state_dict) rrOrQrRrTrVrWrYrZr\r^rrr r!s ) zLearningRateFinder.__init__cCs:|j|jd|j|jd|j|jdS)z9Restores the model and optimizer to their initial states.rOrQN)rOload_state_dictrgretrieverQtorfr$rrr resetszLearningRateFinder.resetg$@dexpg?rA train_loaderr val_loaderDataLoader | Nonerrrstart_lr float | Noneend_lrfloatnum_iter step_modestrsmooth_f diverge_thaccumulation_stepsnon_blocking_transferr3cCsggd|_td }|j|j||r|||dkr%td|dkr3t |j ||}n|dkrAt |j ||}ntd|| dksP| dkrTtd t |||}|rbt |||}|jrstrsttjd d }tjj}nt}t}||D]k}|jrtstd |dd ||j|| | d}|r|j|| d}|jd|d||dkr|}n| dkr| |d| |jdd}||kr|}|jd||| |kr|jr|dnq{| r|jrtd|dSdS)ajPerforms the learning rate range test. Args: train_loader: training set data loader. val_loader: validation data loader (if desired). image_extractor: callable function to get the image from a batch of data. Default: `x["image"] if isinstance(x, dict) else x[0]`. label_extractor: callable function to get the label from a batch of data. Default: `x["label"] if isinstance(x, dict) else x[1]`. start_lr : the starting learning rate for the range test. The default is the optimizer's learning rate. end_lr: the maximum learning rate to test. The test may stop earlier than this if the result starts diverging. num_iter: the max number of iterations for test. step_mode: schedule for increasing learning rate: (`linear` or `exp`). smooth_f: the loss smoothing factor within the `[0, 1[` interval. Disabled if set to `0`, otherwise loss is smoothed using exponential smoothing. diverge_th: test is stopped when loss surpasses threshold: `diverge_th * best_loss`. accumulation_steps: steps for gradient accumulation. If set to `1`, gradients are not accumulated. non_blocking_transfer: when `True`, moves data to device asynchronously if possible, e.g., moving CPU Tensors with pinned memory to CUDA devices. auto_reset: if `True`, returns model and optimizer to original states at end of test. Returns: None r_infrAz `num_iter` must be larger than 1rolinearz#expected one of (exp, linear), got rz$smooth_f is outside the range [0, 1[zComputing optimal learning rate)descz+Computing optimal learning rate, iteration /)r~r`raz%Stopping early, the loss has divergedzResetting model and optimizerN)rdrwrOrlrTrc_set_learning_raterlowerr rQr r2r<r^has_tqdmrrtrangewriterangeprint _train_batch _validateappendget_lrsteprm)rrqrrrrrtrvrxryr{r|r}r~r3 best_lossZ lr_schedule train_iterval_iterrZtprint iterationrarrr range_tests` .            zLearningRateFinder.range_testnew_lrs float | listcCs\t|ts|gt|jj}t|t|jjkrtdt|jj|D]\}}||d<q#dS)z#Set learning rate(s) for optimizer.zYLength of `new_lrs` is not equal to the number of parameter groups in the given optimizerr`N)rlistr>rQ param_groupsrzip)rr param_groupnew_lrrrr rys  z%LearningRateFinder._set_learning_ratecCs"|jjD] }d|vrtdqdS)z/Check optimizer doesn't already have scheduler. initial_lrz0Optimizer already has a scheduler attached to itN)rQr RuntimeError)rrrrr rcs z'LearningRateFinder._check_for_schedulerrr2c Cs|jd}|jt|D]^}t|\}}t||g|j|d\}}||}|||} | |} |j rdt |jdrd|d|dk} t j j j | |j| d } | Wdn1s^wYn| || 7}q|j|S)NrrT non_blockingZ _amp_stashrA) delay_unscale)rOtrainrQ zero_gradrr*rrTrRrYhasattrtorchcudaZ scale_lossbackwarditemr) rrr}r~ total_lossir9r'outputsrarZ scaled_lossrrr rs&        zLearningRateFinder._train_batchrr<cCsd}t|j0|D]%\}}t||g|j|d\}}||}|||}||t|7}q Wdn1s:wY|t|jS)Nrr)r rOrrTrRrr>r#)rrr~Z running_lossr9r'rrarrr rs      zLearningRateFinder._validater skip_startskip_endtuple[list, list]cCsd|dkrtd|dkrtd|jd}|jd}t||d}|||}|||}||fS)zGet learning rates and their corresponding losses Args: skip_start: number of batches to trim from the start. skip_end: number of batches to trim from the end. rzskip_start cannot be negativezskip_end cannot be negativer`rarA)rrdr>)rrrlrslossesend_idxrrr get_lrs_and_lossess    z%LearningRateFinder.get_lrs_and_losses'tuple[float, float] | tuple[None, None]cCsT|||\}}ztt|}||||fWSty)tdYdSw)aCGet learning rate which has steepest gradient and its corresponding loss Args: skip_start: number of batches to trim from the start. skip_end: number of batches to trim from the end. Returns: Learning rate which has steepest gradient and its corresponding loss zBFailed to compute the gradients, there might not be enough points.)NN)rnpgradientarrayargminrr)rrrrrZ min_grad_idxrrr get_steepest_gradients  z(LearningRateFinder.get_steepest_gradientlog_lrax Any | None steepest_lrc Csts tddS|||\}}d}|durt\}}||||rE|||\} } | durE| durE|j| | dddddd| |rL| d | d | d |dur^t |S) aPlots the learning rate range test. Args: skip_start: number of batches to trim from the start. skip_end: number of batches to trim from the start. log_lr: True to plot the learning rate in a logarithmic scale; otherwise, plotted in a linear scale. ax: the plot is created in the specified matplotlib axes object and the figure is not be shown. If `None`, then the figure and axes object are created in this method and the figure is shown. steepest_lr: plot the learning rate which had the steepest gradient. Returns: The `matplotlib.axes.Axes` object that contains the plot. Returns `None` if `matplotlib` is not installed. z(Matplotlib is missing, can't plot resultNKoredzsteepest gradient)smarkercolorzorderrMlogz Learning rateLoss)has_matplotlibwarningswarnrpltsubplotsplotrscatterlegend set_xscale set_xlabel set_ylabelshow) rrrrrrrrfigZlr_at_steepest_gradZloss_at_steepest_gradrrr rs8       zLearningRateFinder.plot)rOrPrQrrRrSrTrUrVr4rWrXrYr4rZr[r\r]r^r4rr)rr)rqr rrrsrrrrrtrurvrwrxr]ryrzr{rwr|r]r}r]r~r4r3r4rr)rrrrr:)rr2r}r]r~r4rrw)rr<r~r4rrw)rr)rr]rr]rr)rr]rr]rr)rrTNT) rr]rr]rr4rrrr4rr)r.r/r0rBpicklerr!rmrKrNrrrcrrrrrrrrr rsJ+ @  z  #  )rCrrrD), __future__rrtypesr functoolsrtypingrrrnumpyrrtorch.nnnn torch.optimrtorch.serializationrtorch.utils.datar monai.networks.utilsr Zmonai.optimizers.lr_schedulerr r monai.utilsr rrmatplotlib.pyplotpyplotrrrr__all__rr2r<rKrNrrrrr s:           %