U PhU@sTddlmZddlZddlZddlZddlmZddlmZm Z m Z ddl Z ddl Z ddlmZddlmZddlmZddlmZddlmZdd lmZmZdd lmZmZmZerddlm Z!d Z"ddl#Z#d Z$ned \Z!Z"ed \Z#Z$dgZ%GdddZ&Gddde&Z'Gddde&Z(dddddZ)dddddZ*GdddZ+dS)) annotationsN)partial) TYPE_CHECKINGAnyCallable) Optimizer)DEFAULT_PROTOCOL) DataLoader) eval_mode) ExponentialLRLinearLR) StateCachercopy_to_deviceoptional_importTzmatplotlib.pyplottqdmLearningRateFinderc@sDeZdZdddddddZeddZd d Zd d Zd dZdS)DataLoaderIterr rNone data_loaderimage_extractorlabel_extractorreturncCs>t|tstdt|d||_t||_||_||_dS)NzLoader has unsupported type: z1. Expected type was `torch.utils.data.DataLoader`) isinstancer ValueErrortyperiter _iteratorrrselfrrrr O/home/dell461/cl/sdc2/HISourceFinder-master-l/src/monai/optimizers/lr_finder.py__init__/s  zDataLoaderIter.__init__cCs|jjSN)rdatasetrr r r!r$9szDataLoaderIter.datasetcCs||}||}||fSr#)rr)r batch_dataimageslabelsr r r!inputs_labels_from_batch=s  z'DataLoaderIter.inputs_labels_from_batchcCs|Sr#r r%r r r!__iter__BszDataLoaderIter.__iter__cCst|j}||Sr#)nextrr))rbatchr r r!__next__Es zDataLoaderIter.__next__N) __name__ __module__ __qualname__r"propertyr$r)r*r-r r r r!r-s   rcs4eZdZd ddddddfdd Zd d ZZS) TrainDataLoaderIterTr rboolr)rrr auto_resetrcst|||||_dSr#)superr"r4)rrrrr4 __class__r r!r"LszTrainDataLoaderIter.__init__cCsfzt|j}||\}}Wn@tk r\|js4t|j|_t|j}||\}}YnX||fSr#)r+rr) StopIterationr4rr)rr,inputsr(r r r!r-Rs   zTrainDataLoaderIter.__next__)T)r.r/r0r"r- __classcell__r r r6r!r2Jsr2cs@eZdZdZdddddfdd Zdd Zfd d ZZS) 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()` ``` r rrrcs&t|||t|j|_d|_dSNr)r5r"lenr run_limit run_counterrr6r r!r"us zValDataLoaderIter.__init__cCs"|j|jkrt|j|_d|_|Sr<)r?r>rrrr%r r r!r*zs  zValDataLoaderIter.__iter__cs|jd7_tS)N)r?r5r-r%r6r r!r-szValDataLoaderIter.__next__)r.r/r0__doc__r"r*r-r:r r r6r!r;`sr;rz torch.Tensor)xrcCst|tr|dn|d}|S)z3Default callable for getting image from batch data.imagerrdictrBoutr r r!default_image_extractorsrHcCst|tr|dn|d}|S)z3Default callable for getting label from batch data.labelr@rDrFr r r!default_label_extractorsrJc@seZdZdZddddeedfddddd d d d d d d d ddZd dddZdee dddddddddf ddddddd ddd d d d d d d!d"Z d#d d$d%d&Z d'd(Z d=d)d d dd*d+d,Z d>d-d dd.d/d0Zd?d d d2d3d4d5Zd@d d d6d3d7d8ZdAd d d d9d d9d:d;d<ZdS)BraLearning rate range test. The learning rate range test increases the learning rate in a pre-training run between two boundaries in a linear or exponential manner. It provides valuable information on how well the network can be trained over a range of learning rates and what is the optimal learning rate. Example (fastai approach): >>> 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 NTFz nn.Modulerztorch.nn.Modulezstr | torch.device | Noner3z str | Noneztypes.ModuleTypeintr) model optimizer criteriondevice memory_cache cache_diramp pickle_modulepickle_protocolverboserc Cs||_|||_||_ggd|_||_||_||_| |_t |j j |_ t |||| d|_|jd|j|jd|j|r|n|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_memoryrQrSrTrLrMN)rM_check_for_schedulerrLrNhistoryrPrQrRrUr+ parametersrO model_devicer state_cacherstore state_dict) rrLrMrNrOrPrQrRrSrTrUr r r!r"s&) zLearningRateFinder.__init__)rcCs:|j|jd|j|jd|j|jdS)z9Restores the model and optimizer to their initial states.rLrMN)rLload_state_dictr^retrieverMtor]r%r r r!resetszLearningRateFinder.resetg$@dexpg?r@r zDataLoader | Nonerz float | Nonefloatstr) train_loader val_loaderrrstart_lrend_lrnum_iter step_modesmooth_f diverge_thaccumulation_stepsnon_blocking_transferr4rcCsggd|_td }|j|j||r:|||dkrJtd|dkrft |j ||}n*|dkrt |j ||}ntd|| dks| dkrtd t |||}|rt |||}|jrtrttjd d }tjj}nt}t}||D]}|jr ts td |dd ||j|| | d}|rD|j|| d}|jd|d||dkrt|}n6| dkr| |d| |jdd}||kr|}|jd||| |kr|jr|dqq| r|jrtd|dS)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 rVinfr@z `num_iter` must be larger than 1rflinearz#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 /)rsrWrXz%Stopping early, the loss has divergedzResetting model and optimizerN)r[rhrLrcrOrZ_set_learning_raterlowerr rMr r2r;rUhas_tqdmrrtrangewriterangeprint _train_batch _validateappendget_lrsteprd)rrjrkrrrlrmrnrorprqrrrsr4 best_lossZ lr_schedule train_iterval_iterr|Ztprint iterationrXr r r! range_tests\.              zLearningRateFinder.range_testz float | list)new_lrsrcCs\t|ts|gt|jj}t|t|jjkr8tdt|jj|D]\}}||d<qFdS)z#Set learning rate(s) for optimizer.zYLength of `new_lrs` is not equal to the number of parameter groups in the given optimizerrWN)rlistr=rM param_groupsrzip)rr param_groupnew_lrr r r!ryys z%LearningRateFinder._set_learning_ratecCs"|jjD]}d|krtdqdS)z/Check optimizer doesn't already have scheduler. initial_lrz0Optimizer already has a scheduler attached to itN)rMr RuntimeError)rrr r r!rZs z'LearningRateFinder._check_for_schedulerr2)rrrrsrc Cs|jd}|jt|D]}t|\}}t||g|j|d\}}||}|||} | |} |j rt |jdr|d|dk} t j j j | |j| d} | W5QRXn| || 7}q |j|S)NrrO non_blockingZ _amp_stashr@) delay_unscale)rLtrainrM zero_gradr~r+rrOrNrRhasattrtorchcudaZ scale_lossbackwarditemr) rrrrrs total_lossir9r(outputsrXrZ scaled_lossr r r!rs"       zLearningRateFinder._train_batchr;)rrsrc Csxd}t|jV|D]J\}}t||g|j|d\}}||}|||}||t|7}qW5QRX|t|jS)Nrr)r rLrrOrNrr=r$)rrrsZ running_lossr9r(rrXr r r!rs      zLearningRateFinder._validaterztuple[list, list]) skip_startskip_endrcCsd|dkrtd|dkr td|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 negativerWrXr@)rr[r=)rrrlrslossesend_idxr r r!get_lrs_and_lossess    z%LearningRateFinder.get_lrs_and_lossesz)tuple[float, float] | tuple[(None, None)]cCsZ|||\}}z&tt|}||||fWStk rTtdYdSXdS)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.)NNN)rnpgradientarrayargminrr)rrrrrZ min_grad_idxr r r!get_steepest_gradients z(LearningRateFinder.get_steepest_gradientz Any | None)rrlog_lrax steepest_lrrc CststddS|||\}}d}|dkr:t\}}||||r|||\} } | dk r|j| | dddddd| |r| d | d | d |dk rt |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)smarkercolorzorderrIlogz Learning rateLoss)has_matplotlibwarningswarnrpltsubplotsplotrscatterlegend set_xscale set_xlabel set_ylabelshow) rrrrrrrrfigZlr_at_steepest_gradZloss_at_steepest_gradr r r!rs8       zLearningRateFinder.plot)T)T)rr)rr)rrTNT)r.r/r0rApicklerr"rdrHrJrryrZrrrrrr r r r!rsH+$@ *z #), __future__rrtypesr functoolsrtypingrrrnumpyrrtorch.nnnn torch.optimrtorch.serializationrtorch.utils.datar monai.networks.utilsr Zmonai.optimizers.lr_schedulerr r monai.utilsr rrmatplotlib.pyplotpyplotrrrr{__all__rr2r;rHrJrr r r r! s8          %