o %ݫi@sdZddlZddlZddlmZddlmZddlmZeeZ d(ddZ ej Gdd d Z ej Gd d d Z Gd d d Zej GdddZGdddZej GdddZej GdddZej GdddZej GdddZej GdddZej GdddZej GdddZej Gd d!d!Zej Gd"d#d#ZGd$d%d%ejZej Gd&d'd'ZdS))z Schedulers for updating hyperparameters (such as learning rate). Authors * Mirco Ravanelli 2020 * Peter Plantinga 2020 * Loren Lugosch 2020 * Shucong Zhang 2023 N)nn) checkpoints) get_loggercCsn|dur tt|j}n|}|D]$}|j|d}||kr4||j|d<||j|d<td||fqdS)aChange the learning rate value within an optimizer. Arguments --------- optimizer : torch.optim object Updates the learning rate for this optimizer. new_lr : float The new value to use for the learning rate. param_group : list of int The param group indices to update. If not provided, all groups updated. Example ------- >>> from torch.optim import SGD >>> from speechbrain.nnet.linear import Linear >>> model = Linear(n_neurons=10, input_size=10) >>> optimizer = SGD(model.parameters(), lr=0.1) >>> update_learning_rate(optimizer, 0.2) >>> optimizer.param_groups[0]["lr"] 0.2 Nlrprev_lrzChanging lr from %.2g to %.2g)rangelen param_groupsloggerinfo) optimizernew_lr param_groupgroupsiold_lrrO/home/ubuntu/.local/lib/python3.10/site-packages/speechbrain/nnet/schedulers.pyupdate_learning_ratesrcsHeZdZdZd fdd ZddZejddZej dd d Z Z S)WarmAndExpDecayLRScheduleaUWarms up linearly, and then decay exponentially to ('lr' / 'decay_factor') in 'total_steps' steps. Arguments --------- lr : float The max learning rate to reach after warmup. n_warmup_steps : int Number of warmup steps (following a linear increase). total_steps : int Total number of steps (used to decay). decay_factor : float Decay factor applied every decay_every steps. (default: 0.01) Example ------- >>> from speechbrain.nnet.linear import Linear >>> inp_tensor = torch.rand([1,660,3]) >>> model = Linear(input_size=3, n_neurons=4) >>> optim = torch.optim.Adam(model.parameters(), lr=1) >>> output = model(inp_tensor) >>> scheduler = WarmAndExpDecayLRSchedule(lr=1, n_warmup_steps=2, decay_factor=0.01, total_steps=6) >>> scheduler(optim) >>> optim.param_groups[0]["lr"] 0.0 >>> scheduler(optim) >>> optim.param_groups[0]["lr"] 0.5 >>> scheduler(optim) >>> optim.param_groups[0]["lr"] 1 >>> scheduler(optim) >>> optim.param_groups[0]["lr"] 0.31622776601683794 皙?cs<tt|||_d|_||_||_||j|_d|_dSNr) superr__init__base_lr current_lrn_warmup_steps decay_factor decay_steps current_step)selfrr total_stepsr __class__rrrbs  z"WarmAndExpDecayLRSchedule.__init__cCsv|j|jkr|j|j|j}n|j|j|j|j|j}t|j|}|jD]}||d<q(||_|jd7_dS)Nr)rrrrrminr r)r optr decayed_lrrrrr__call__ks    z"WarmAndExpDecayLRSchedule.__call__cC*|j|j|j|j|jd}t||dS)0Saves the current metrics on the specified path.)rrrrrN)rrrrrtorchsaver pathdatarrrr,{zWarmAndExpDecayLRSchedule.saveFNcCsD~~t|}|d|_|d|_|d|_|d|_|d|_dS)Loads the needed information.rrrrrN)r+loadrrrrrr r. end_of_epochdevicer/rrrr2s     zWarmAndExpDecayLRSchedule.load)rFN __name__ __module__ __qualname____doc__rr(r mark_as_saverr,mark_as_loaderr2 __classcell__rrr"rr<s$   rc@sFeZdZdZ   dddZddZejd d Zej dd d Z dS)NewBobScheduleraScheduler with new-bob technique, used for LR annealing. The learning rate is annealed based on the validation performance. In particular: if (past_loss-current_loss)/past_loss< impr_threshold: lr=lr * annealing_factor. Arguments --------- initial_value : float The initial hyperparameter value. annealing_factor : float It is annealing factor used in new_bob strategy. improvement_threshold : float It is the improvement rate between losses used to perform learning annealing in new_bob strategy. patient : int When the annealing condition is violated patient times, the learning rate is finally reduced. Example ------- >>> scheduler = NewBobScheduler(initial_value=1.0) >>> scheduler(metric_value=10.0) (1.0, 1.0) >>> scheduler(metric_value=2.0) (1.0, 1.0) >>> scheduler(metric_value=2.5) (1.0, 0.5) ?{Gzd?rcCs*||_||_||_||_g|_|j|_dSN)hyperparam_valueannealing_factorimprovement_thresholdpatient metric_valuescurrent_patient)r initial_valuerDrErFrrrrs  zNewBobScheduler.__init__cCs|j}}t|jdkr9|jd}|dkrd}n|||}||jkr9|jdkr2||j9}|j|_n|jd8_|j|||_||fS)aReturns the current and new value for the hyperparameter. Arguments --------- metric_value : int A number for determining whether to change the hyperparameter value. Returns ------- Current and new hyperparam value. rr$)rCrrGrErHrDrFappend)r metric_value old_value new_value prev_metric improvementrrrr(s        zNewBobScheduler.__call__cC"|j|j|jd}t||dS)r*)rCrGrHN)rCrGrHr+r,r-rrrr, zNewBobScheduler.saveFcC.~t|}|d|_|d|_|d|_dS)r1rCrGrHN)r+r2rCrGrHr r.r4r/rrrr2    zNewBobScheduler.loadN)r@rArF r8r9r:r;rr(rr<r,r=r2rrrrr?s!   r?c@s eZdZdZddZddZdS)LinearScheduleraScheduler with linear annealing technique. The learning rate linearly decays over the specified number of epochs. Arguments --------- initial_value : float The value upon initialization. final_value : float The value used when the epoch count reaches ``epoch_count - 1``. epoch_count : int Number of epochs. Example ------- >>> scheduler = LinearScheduler(1.0, 0.0, 4) >>> scheduler(current_epoch=1) (1.0, 0.666...) >>> scheduler(current_epoch=2) (0.666..., 0.333...) >>> scheduler(current_epoch=3) (0.333..., 0.0) >>> scheduler(current_epoch=4) (0.0, 0.0) cCstj|||d|_dS)Nsteps)r+linspacetolistvalue_at_epoch)r rI final_value epoch_countrrrrs zLinearScheduler.__init__cCs6td|d}t|t|jd}|j||j|fS)a Returns the current and new value for the hyperparameter. Arguments --------- current_epoch : int Number of times the dataset has been iterated. Returns ------- Current and new hyperparam value. rr$)maxr%rr])r current_epoch old_indexindexrrrr(s zLinearScheduler.__call__N)r8r9r:r;rr(rrrrrXs rXc@sFeZdZdZddZddZddZejdd Z ej dd d Z d S)LinearWarmupScheduleraLCreate a schedule with a learning rate that decreases linearly from the initial lr set in the optimizer to 0, after a warmup period during which it increases linearly from 0 to the initial lr set in the optimizer. * Ge Li 2022 Arguments --------- initial_value : float The value upon initialization (lr0). num_warmup_steps : int Number of warmup steps. The learning rate reaches lr0 at ``num_warmup_steps + 1`` step. num_training_steps : int The total number of training steps. Example ------- >>> scheduler = LinearWarmupScheduler(1.0, 2, 4) >>> scheduler.get_next_value() 0.0 >>> scheduler.get_next_value() 0.5 >>> scheduler.get_next_value() 1.0 >>> scheduler.get_next_value() 0.5 >>> scheduler.get_next_value() 0.0 cCs||_||_||_d|_dSr)lr0num_warmup_stepsnum_training_stepsr)r rIrfrgrrrrHs zLinearWarmupScheduler.__init__c CsX||jkrt|ttd|j|jS|jtdt|j|ttd|j|jS)aReturns the current and new value for the hyperparameter. Arguments --------- current_step : int Number of steps the model has been updated. Returns ------- Current and new hyperparam value. r$)rffloatr`rerg)r rrrr calculate_lrNs  z"LinearWarmupScheduler.calculate_lrcCs||j}|jd7_|S)z>> scheduler = StepScheduler(initial_value=1.0) >>> scheduler(current_epoch=1) (1.0, 0.5) >>> scheduler(current_epoch=2) (0.5, 0.5) >>> scheduler(current_epoch=3) (0.5, 0.25) r@NcCsL||_|r|s |r td|||_d|_dS|p|j|_|p"|j|_dS)NzBhalf_life cannot be used together with decay_factor and decay_drop?)rI ValueError_compute_half_life_decay_factorr decay_dropDEFAULT_DECAY_FACTORDEFAULT_DECAY_DROP)r rIrrq half_liferrrrs   zStepScheduler.__init__cCsttd |S)Nrm)mathexplog)r rtrrrrpsz-StepScheduler._compute_half_life_decay_factorcCs ||d}||}||fS)zReturns current and new hyperparameter value. Arguments --------- current_epoch : int Number of times the dataset has been iterated. Returns ------- Current and new hyperparam value. r$)_compute_value)r ra current_value next_valuerrrr(s zStepScheduler.__call__cCs$|jt|jtd||jSNr$)rIrupowrfloorrq)r rarrrrxszStepScheduler._compute_valueNNN) r8r9r:r;rrrsrrpr(rxrrrrrls   rlc@sHeZdZdZdddZddZddZejd d Z ej dd d Z dS) NoamScheduleraCThe is an implementation of the transformer's learning rate scheduler with warmup. Reference: https://arxiv.org/abs/1706.03762 Note: this scheduler anneals the lr at each update of the model's weight, and n_steps must be saved for restarting. Arguments --------- lr_initial : float Initial learning rate (i.e. the lr used at epoch 0). n_warmup_steps : int number of warm-up steps model_size : int size of transformer embed_dim. It is used to scale the maximum learning rate value reached by the scheduler. It is divided by model_size ** (0.5). If not specified the maximum learning rate value is instead multiplied by warmup_steps ** (0.5). Example ------- >>> from speechbrain.nnet.linear import Linear >>> inp_tensor = torch.rand([1,660,3]) >>> model = Linear(input_size=3, n_neurons=4) >>> optim = torch.optim.Adam(model.parameters(), lr=1) >>> output = model(inp_tensor) >>> scheduler =NoamScheduler(optim.param_groups[0]["lr"], 3) >>> curr_lr,next_lr=scheduler(optim) >>> optim.param_groups[0]["lr"] 0.3333333333333333 >>> curr_lr,next_lr=scheduler(optim) >>> optim.param_groups[0]["lr"] 0.6666666666666666 >>> curr_lr,next_lr=scheduler(optim) >>> optim.param_groups[0]["lr"] 0.9999999999999999 NcCsB||_||_||_g|_d|_|d|_|dur|d|_dSdSNrr@) lr_initialrrlossesn_steps normalize)r rr model_sizerrrrs zNoamScheduler.__init__cCL|jd7_|jdd}|j|}|jD]}||d<q||_||fSa/ Arguments --------- opt : optimizer The optimizer to update using this scheduler. Returns ------- current_lr : float The learning rate before the update. lr : float The learning rate after the update. r$rrrr r _get_lr_scalerr r&rrrrrrr(  zNoamScheduler.__call__cCs*|j|j}}|jt|d||dSNr)rrrr%r rrrrrrszNoamScheduler._get_lr_scalecC|j|jd}t||dSr*)rrNrrr+r,r-rrrr,!zNoamScheduler.saveFcC$~t|}|d|_|d|_dSr1rrNr+r2rrrTrrrr2'  zNoamScheduler.loadrBrV r8r9r:r;rr(rrr<r,r=r2rrrrrs $  rc@sJeZdZdZ dddZddZddZejd d Z ej dd d Z dS)NoamIntervalScheduleraA combination of Noam Scheduler and Interval Scheduler. The scheduler behaves as a Noam Scheduler, and anneals the learning rate at designed steps with designed decays. Note: this scheduler anneals the lr at each update of the model's weight, and n_steps must be saved for restarting. Arguments --------- lr_initial : float Initial learning rate (i.e. the lr used at epoch 0). n_warmup_steps : int number of warm-up steps. anneal_steps: list Pre-designed steps where the learning rate is to be annealed. anneal_rates: list Pre-designed decay rate for each anneal step. model_size : int size of transformer embed_dim. It is used to scale the maximum learning rate value reached by the scheduler. It is divided by model_size ** (0.5). If not specified the maximum learning rate value is instead multiplied by warmup_steps ** (0.5). Example ------- >>> from speechbrain.nnet.linear import Linear >>> inp_tensor = torch.rand([1,660,3]) >>> model = Linear(input_size=3, n_neurons=4) >>> optim = torch.optim.Adam(model.parameters(), lr=1) >>> output = model(inp_tensor) >>> scheduler = NoamIntervalScheduler( ... lr_initial=optim.param_groups[0]["lr"], ... n_warmup_steps=3, ... anneal_steps=[6, 9], ... anneal_rates=[0.5, 0.1], ... ) >>> for _ in range(10): ... curr_lr,next_lr=scheduler(optim) ... print(optim.param_groups[0]["lr"]) 0.3333333333333333 0.6666666666666666 0.9999999999999999 0.8660254037844386 0.7745966692414833 0.7071067811865475 0.3273268353539886 0.3061862178478973 0.28867513459481287 0.027386127875258306 NcCsN||_||_||_g|_d|_|d|_||_||_|dur%|d|_dSdSr)rrrrrr anneal_steps anneal_rates)r rrrrrrrrrds zNoamIntervalScheduler.__init__cCrrrrrrrr(wrzNoamIntervalScheduler.__call__cCs`|j|j}}|jt|d||d}tt|jD]}|j|j|kr-||j|}q|Sr)rrrr%rrrr)r rrlr_scalerrrrrsz#NoamIntervalScheduler._get_lr_scalecCrrrr-rrrr,rzNoamIntervalScheduler.saveFcC&~~t|}|d|_|d|_dSrrr3rrrr2   zNoamIntervalScheduler.loadrBr6rrrrrr0s8  rc@sFeZdZdZddZddZddZejdd Z ej dd d Z d S)LinearNoamSchedulera!The is an implementation of the extended Noam scheduler in the Squeezeformer paper. Reference: https://arxiv.org/pdf/2206.00888.pdf Note: this scheduler anneals the lr at each update of the model's weight, and n_steps must be saved for restarting. Arguments --------- lr_initial : float Initial learning rate (i.e. the lr used at epoch 0). n_warmup_steps : int number of warm-up steps. n_keep_steps : int after warmp-up steps, number of steps that the lr is kept unchanged. Example ------- >>> from speechbrain.nnet.linear import Linear >>> inp_tensor = torch.rand([1,660,3]) >>> model = Linear(input_size=3, n_neurons=4) >>> optim = torch.optim.Adam(model.parameters(), lr=1) >>> output = model(inp_tensor) >>> scheduler =LinearNoamScheduler(optim.param_groups[0]["lr"], 2, 2) >>> curr_lr,next_lr=scheduler(optim) >>> optim.param_groups[0]["lr"] 0.5 >>> curr_lr,next_lr=scheduler(optim) >>> optim.param_groups[0]["lr"] 1.0 >>> curr_lr,next_lr=scheduler(optim) >>> optim.param_groups[0]["lr"] 1.0 >>> curr_lr,next_lr=scheduler(optim) >>> optim.param_groups[0]["lr"] 1.0 >>> curr_lr,next_lr=scheduler(optim) >>> optim.param_groups[0]["lr"] 0.6666666666666666 cCs(||_||_||_||_g|_d|_dSr)rr n_keep_stepsrrr)r rrrrrrrs  zLinearNoamScheduler.__init__cCrrrrrrrr(rzLinearNoamScheduler.__call__cCsB|j|j}}||kr|d|S||j|krdS|||jS)Nrhrn)rrrrrrrrs  z!LinearNoamScheduler._get_lr_scalecCrrrr-rrrr,rzLinearNoamScheduler.saveFNcCrrrr3rrrr2rzLinearNoamScheduler.loadr6rrrrrrs( rc@sHeZdZdZdddZddZdd Zejd d Z ej dd dZ dS)CyclicCosineScheduleracThe is an implementation of the Cyclic-Cosine learning rate scheduler with warmup. Reference: https://openreview.net/pdf?id=BJYwwY9ll Note: this scheduler anneals the lr at each update of the model's weight, and n_steps must be saved for restarting. Arguments --------- n_warmup_steps : int Number of warm up steps. lr_initial : float Initial learning rate (i.e. the lr used at epoch 0). total_steps : int Total number of updating steps. Example ------- >>> from speechbrain.nnet.linear import Linear >>> inp_tensor = torch.rand([1,660,3]) >>> model = Linear(input_size=3, n_neurons=4) >>> optim = torch.optim.Adam(model.parameters(), lr=1) >>> output = model(inp_tensor) >>> scheduler =CyclicCosineScheduler(3, optim.param_groups[0]["lr"]) >>> curr_lr,next_lr=scheduler(optim) >>> optim.param_groups[0]["lr"] 0.9999999990130395 >>> curr_lr,next_lr=scheduler(optim) >>> optim.param_groups[0]["lr"] 0.9999999997532598 >>> curr_lr,next_lr=scheduler(optim) >>> optim.param_groups[0]["lr"] 1.0 N順cCs:||_g|_||_||_||_d|_d||d|_dS)Nrr$r)rr initial_lrrtotalrr)r rrr!rrrr7szCyclicCosineScheduler.__init__cCs\|jd7_|jdur|jdd}n|j}||}|jD]}||d<q ||_||fS)a9 Arguments --------- opt : list of optimizers The optimizers to update using this scheduler. Returns ------- current_lr : float The learning rate before the update. lr : float The learning rate after the update. r$Nrr)rrr rrrrrrr(As    zCyclicCosineScheduler.__call__cCs0|j|j}}dttj|||jdS)Nr@r$)rrrucospirrrrrr_sz#CyclicCosineScheduler._get_lr_scalecCrrrr-rrrr,erzCyclicCosineScheduler.saveFcCrrrrTrrrr2krzCyclicCosineScheduler.load)NrrVrrrrrrs #  rc@sBeZdZdZ dddZdd Zejd d Zej dd dZ dS)ReduceLROnPlateaualLearning rate scheduler which decreases the learning rate if the loss function of interest gets stuck on a plateau, or starts to increase. The difference from NewBobLRScheduler is that, this one keeps a memory of the last step where do not observe improvement, and compares against that particular loss value as opposed to the most recent loss. Arguments --------- lr_min : float The minimum allowable learning rate. factor : float Factor with which to reduce the learning rate. patience : int How many epochs to wait before reducing the learning rate. dont_halve_until_epoch : int Number of epochs to wait until halving. Example ------- >>> from torch.optim import Adam >>> from speechbrain.nnet.linear import Linear >>> inp_tensor = torch.rand([1,660,3]) >>> model = Linear(n_neurons=10, input_size=3) >>> optim = Adam(lr=1.0, params=model.parameters()) >>> output = model(inp_tensor) >>> scheduler = ReduceLROnPlateau(0.25, 0.5, 2, 1) >>> curr_lr,next_lr=scheduler([optim],current_epoch=1, current_loss=10.0) >>> curr_lr,next_lr=scheduler([optim],current_epoch=2, current_loss=11.0) >>> curr_lr,next_lr=scheduler([optim],current_epoch=3, current_loss=13.0) >>> curr_lr,next_lr=scheduler([optim],current_epoch=4, current_loss=14.0) >>> next_lr 0.5 :0yE>r@rmAcCs.||_||_||_d|_g|_||_d|_dS)Nri)lr_minfactorpatiencepatience_counterrdont_halve_until_epochanchor)r rrrrrrrrs zReduceLROnPlateau.__init__cCs|D]D}|jdd}||jkr|}||_n*||jkr$d|_|}||_n||jkr8|j|jkr8|jd|_|}n||j}d|_t||j}q|j |||fS)a Arguments --------- optim_list : list of optimizers The optimizers to update using this scheduler. current_epoch : int Number of times the dataset has been iterated. current_loss : int A number for determining whether to change the learning rate. Returns ------- current_lr : float The learning rate before the update. next_lr : float The learning rate after the update. rrr$) r rrrrrr`rrrK)r optim_listra current_lossr&rnext_lrrrrr(s$       zReduceLROnPlateau.__call__cCrQ)r*)rrrN)rrrr+r,r-rrrr,rRzReduceLROnPlateau.saveFcCrS)r1rrrN)r+r2rrrrTrrrr2rUzReduceLROnPlateau.loadN)rr@rmrrVrWrrrrrts#  /  rcspeZdZdZ       dfd d Zdd d Zd dZddZddZe j ddZ e j dddZ ZS)CyclicLRSchedulera` This implements a cyclical learning rate policy (CLR). The method cycles the learning rate between two boundaries with some constant frequency, as detailed in this paper (https://arxiv.org/abs/1506.01186). The amplitude of the cycle can be scaled on a per-iteration or per-cycle basis. This class has three built-in policies, as put forth in the paper. "triangular": A basic triangular cycle w/ no amplitude scaling. "triangular2": A basic triangular cycle that scales initial amplitude by half each cycle. "exp_range": A cycle that scales initial amplitude by gamma**(cycle iterations) at each cycle iteration. For more detail, please see the reference paper. Arguments --------- base_lr : float initial learning rate which is the lower boundary in the cycle. max_lr : float upper boundary in the cycle. Functionally, it defines the cycle amplitude (max_lr - base_lr). The lr at any cycle is the sum of base_lr and some scaling of the amplitude; therefore max_lr may not actually be reached depending on scaling function. step_size : int number of training iterations per half cycle. The authors suggest setting step_size 2-8 x training iterations in epoch. mode : str one of {triangular, triangular2, exp_range}. Default 'triangular'. Values correspond to policies detailed above. If scale_fn is not None, this argument is ignored. gamma : float constant in 'exp_range' scaling function: gamma**(cycle iterations) scale_fn : lambda function Custom scaling policy defined by a single argument lambda function, where 0 <= scale_fn(x) <= 1 for all x >= 0. mode parameter is ignored scale_mode : str {'cycle', 'iterations'}. Defines whether scale_fn is evaluated on cycle number or cycle iterations (training iterations since start of cycle). Default is 'cycle'. Example ------- >>> from speechbrain.nnet.linear import Linear >>> inp_tensor = torch.rand([1,660,3]) >>> model = Linear(input_size=3, n_neurons=4) >>> optim = torch.optim.Adam(model.parameters(), lr=1) >>> output = model(inp_tensor) >>> scheduler = CyclicLRScheduler(base_lr=0.1, max_lr=0.3, step_size=2) >>> scheduler.on_batch_end(optim) >>> optim.param_groups[0]["lr"] 0.2 >>> scheduler.on_batch_end(optim) >>> optim.param_groups[0]["lr"] 0.3 >>> scheduler.on_batch_end(optim) >>> optim.param_groups[0]["lr"] 0.2 MbP?~jtx?@@ triangularrnNcyclecstg|_||_||_||_||_|_|durG|jdkr)dd|_d|_ n$|jdkr7dd|_d|_ n|jdkrFfdd|_d |_ n||_||_ d |_ | dS) NrcSsdS)NrnrxrrrBsz,CyclicLRScheduler.__init__..r triangular2cSsdd|dS)Nr$g@rrrrrrEs exp_rangecs|SrBrrgammarrrHs iterationsrh) rrrrmax_lr step_sizemoderscale_fn scale_modeclr_iterations_reset)r rrrrrrrr"rrr.s,       zCyclicLRScheduler.__init__cCs4|dur||_|dur||_|dur||_d|_dS)zQResets cycle iterations. Optional boundary/step size adjustment. Nrh)rrrr)r new_base_lr new_max_lr new_step_sizerrrrQs zCyclicLRScheduler._resetcCs|j}||jd}||fSr{)rclrr)r epochrr rrrr(]szCyclicLRScheduler.__call__cCstd|d|j}t||jd|d}|jdkr3|j|j|jtdd|||S|j|j|jtdd|||S)zClears iterations.r$rmrr) rur}rabsrrrr`r)r rrrrrrrcs zCyclicLRScheduler.clrcCsF|jd7_||j}|jdd}|jD]}||d<q||_dS)z Arguments --------- opt : optimizers The optimizers to update using this scheduler. r$rrN)rrr r)r r&rrrrrr on_batch_endps     zCyclicLRScheduler.on_batch_endcCr)r*)rrN)rrr+r,r-rrrr,rzCyclicLRScheduler.saveFcCr)r1rrN)r+r2rrrTrrrr2rzCyclicLRScheduler.load)rrrrrnNrr~rV)r8r9r:r;rrr(rrrr<r,r=r2r>rrr"rrs$H #   rc@sNeZdZdZddZddZddZdd Zej d d Z ej dd dZ dS)IntervalScheduleraA simple scheduler implementation that sets the learning rate to specific values after a specific number of steps has been reached. Arguments --------- intervals : list a list of dictionaries: {"steps": , "lr": the learning rate} 'steps' indicates the global step count at which a given rate will apply Example ------- >>> import torch >>> from speechbrain.nnet.schedulers import IntervalScheduler >>> from speechbrain.nnet.linear import Linear >>> model = Linear(input_size=3, n_neurons=4) >>> optim = torch.optim.Adam(model.parameters(), lr=1) >>> scheduler = IntervalScheduler( ... intervals=[ ... {"steps": 2, "lr": 0.01}, ... {"steps": 5, "lr": 0.005}, ... {"steps": 9, "lr": 0.001} ... ] ... ) >>> optim.param_groups[0]["lr"] 1 >>> for _ in range(10): ... pre, post = scheduler(optim) ... print(f"{pre} -> {post}") 1 -> 1 1 -> 0.01 0.01 -> 0.01 0.01 -> 0.01 0.01 -> 0.005 0.005 -> 0.005 0.005 -> 0.005 0.005 -> 0.005 0.005 -> 0.001 0.001 -> 0.001 cCs||_d|_g|_|dSr) intervalsrr _compute_next)r rrrrrs zIntervalScheduler.__init__cCsH|jd7_|jdd}||}|jD]}||d<q||_||fSr)rr _get_lrrrrrrr(s   zIntervalScheduler.__call__csfddjD_dS)Ncsg|] }|djkr|qSrY)r).0intervalr rr s z3IntervalScheduler._compute_next..)r_next_intervalsrrrrrs  zIntervalScheduler._compute_nextcCs6|}|jr|jd}|j|dkr|d}|jd=|S)NrrZr)rr)r rr next_intervalrrrrs zIntervalScheduler._get_lrcCrrrr-rrrr,rzIntervalScheduler.saveFcCs,~t|}|d|_|d|_|dSr)r+r2rrrrTrrrr2s     zIntervalScheduler.loadNrV) r8r9r:r;rr(rrrr<r,r=r2rrrrrs) rc@s6eZdZdZddZddZddZejdd Z d S) InverseSquareRootSchedulerzThe Inverse Square Root Scheduler, as defined in the T5 paper https://arxiv.org/pdf/1910.10683.pdf Arguments --------- warmup_steps : int The number of steps over which the learning rate will be constant cCs||_d|_dSr) warmup_stepsr)r rrrrrs z#InverseSquareRootScheduler.__init__cCsF|jd7_|jdd}|}|jD]}||d<q||_||fS)zReturns current and new hyperparameter value. Arguments --------- opt : optimizer The optimizer to update using this scheduler. Returns ------- current and new hyperparam value r$rr)rr rxrrrrrr( s   z#InverseSquareRootScheduler.__call__cCsdtt|j|jSr{)rusqrtr`rrrrrrrx%sz)InverseSquareRootScheduler._compute_valuecCd|ji}t||dS)r*rNrr+r,r-rrrr,( zInverseSquareRootScheduler.saveN) r8r9r:r;rr(rxrr<r,rrrrrs rcsLeZdZdZ  d fdd ZddZejdd Zej dd d Z Z S)WarmCoolDecayLRScheduleaCWarms up linearly, very slowly decays and cools down linearly again at the end of training. This is a three steps scheduler. Reference --------- Scaling Vision Transformers arxiv.org/abs/2106.04560 Arguments --------- lr : float The max learning rate to reach after warmup. warmup : int Number of warmup steps (following a linear increase). cooldown : int Number of cooldown steps (following a linear decrease). total_steps : int Total number of steps (used to decay). decay_factor : float Decay factor applied every decay_every steps. decay_every : int Apply the decay factor to the learning rate every decay_every steps. Example ------- >>> from speechbrain.nnet.linear import Linear >>> inp_tensor = torch.rand([1,660,3]) >>> model = Linear(input_size=3, n_neurons=4) >>> optim = torch.optim.Adam(model.parameters(), lr=1) >>> output = model(inp_tensor) >>> scheduler = WarmCoolDecayLRSchedule(lr=1, warmup=2, total_steps=6, decay_factor=0.5, decay_every=1, cooldown=1) >>> optim.param_groups[0]["lr"] 1 >>> scheduler(optim, 1) >>> optim.param_groups[0]["lr"] 0.5 >>> scheduler(optim, 2) >>> optim.param_groups[0]["lr"] 1.0 >>> scheduler(optim, 3) >>> optim.param_groups[0]["lr"] 0.5 >>> scheduler(optim, 4) >>> optim.param_groups[0]["lr"] 0.25 >>> scheduler(optim, 5) >>> optim.param_groups[0]["lr"] 0.12500000000000003 >>> scheduler(optim, 6) >>> optim.param_groups[0]["lr"] 0.0 ?rcs6t||_||_||_||_t|||_dSrB) rrrwarmupcooldownr!rurwpower)r rrrr!r decay_everyr"rrrfs z WarmCoolDecayLRSchedule.__init__cCs||jkr|j||j}n9||j|jkr9|jt|j|j|j}||j}||j|j}|||}n|jt|j||j}|jD]}||d<qJdS)Nr)rrr!rrurvrr )r r& num_updatesrrdecreasenrrrrr(vs    z WarmCoolDecayLRSchedule.__call__cCr))r*)rrrrr!N)rrrrr!r+r,r-rrrr,r0zWarmCoolDecayLRSchedule.saveFcCsB~t|}|d|_|d|_|d|_|d|_|d|_dS)r1rrrrr!N)r+r2rrrrr!rTrrrr2s     zWarmCoolDecayLRSchedule.load)rrrVr7rrr"rr/s;  rcsNeZdZdZfddZddZejddZej dd d Z d d Z Z S) ScheduledLossaTA convenience class for switching to a different loss function on a schedule Arguments --------- schedule : list a list of dictionaries with the following keys loss_fn: the loss function to use steps: the number of steps to apply before switching to the next one Example ------- >>> loss_fn = ScheduledLoss( ... schedule=[ ... {"steps": 3, "loss_fn": nn.MSELoss()}, ... {"steps": 2, "loss_fn": nn.L1Loss()}, ... {"loss_fn": nn.SmoothL1Loss()} ... ] ... ) >>> x = torch.tensor([1., 2.]) >>> y = torch.tensor([1.5, 2.5]) >>> for idx in range(10): ... loss = loss_fn(x, y) ... print(loss.item()) 0.25 0.25 0.25 0.5 0.5 0.125 0.125 0.125 0.125 0.125 csLtt|s tdtdd|Drtd||_d|_|dS)Nz&At least one schedule item is requiredcss"|] }t|ds|VqdS)loss_fnN)callableget)ritemrrr s z)ScheduledLoss.__init__..z*Each schedule item needs to have at least r)rranyroschedulerfind_next_switch)r rr"rrrs  zScheduledLoss.__init__cOs2|j|jkr ||jd7_|j|i|S)aEComputes the loss at the specified step number. Arguments --------- *args : tuple **kwargs : dict Any arguments passed to this will be passed on to the specified loss_fn Returns ------- result : torch.Tensor the loss value r$)r next_switchrcurrent_loss_fn)r argskwargsrrrforwards zScheduledLoss.forwardcCr)z.Saves the current state on the specified path.rNrr-rrrr,rzScheduledLoss.saveFNcCs t|}|d|_|dS)r1rN)r+r2rrr3rrrr2s   zScheduledLoss.loadcCsJd}|jD]}|dtj}||7}||jkr"|d|_||_dSqdS)zUFinds the threshold at which the next switch will occur based on the schedulerrZrN)rrr+infrrr)r cumulative_stepsr item_stepsrrrrs   zScheduledLoss.find_next_switchr6) r8r9r:r;rrrr<r,r=r2rr>rrr"rrs %   rcsLeZdZdZ  dfdd ZddZejdd Zej dd d Z Z S)TriStageLRScheduleaaWarms up linearly, very slowly decays and cools down linearly again at the end of training. This is a three steps scheduler. Reference https://arxiv.org/pdf/1904.08779.pdf Arguments --------- lr : float The max learning rate to reach after warmup. warmup_steps : int Number of warmup steps (following a linear increase). hold_steps : int Number of holding steps (lr remains unchanged). decay_steps : int Number of decay steps. total_steps : int Total number of steps (used to decay). init_lr_scale : float The initial learning rate scale during warmup phase. final_lr_scale : float The final learning rate scale. Example ------- >>> from speechbrain.nnet.linear import Linear >>> inp_tensor = torch.rand([1,660,3]) >>> model = Linear(input_size=3, n_neurons=4) >>> optim = torch.optim.Adam(model.parameters(), lr=1) >>> output = model(inp_tensor) >>> scheduler = TriStageLRSchedule(lr=1, warmup_steps=2, hold_steps=2, decay_steps=2, total_steps=6, init_lr_scale=0.01, final_lr_scale=0.05) >>> optim.param_groups[0]["lr"] 1 >>> scheduler(optim, 1) >>> optim.param_groups[0]["lr"] 0.505 >>> scheduler(optim, 2) >>> optim.param_groups[0]["lr"] 1 >>> scheduler(optim, 3) >>> optim.param_groups[0]["lr"] 1 >>> scheduler(optim, 4) >>> optim.param_groups[0]["lr"] 1.0 >>> scheduler(optim, 5) >>> optim.param_groups[0]["lr"] 0.223606797749979 >>> scheduler(optim, 6) >>> optim.param_groups[0]["lr"] 0.05000000000000001 {Gz?皙?csttt|||_||_||_||_||_||_||_ |j|j|_ |j|j |j|_ t |j  |j|_dSrB)rrrpeak_lrr hold_stepsrr! init_lr_scalefinal_lr_scaleinit_lr warmup_raterurwr)r rrrrr!rrr"rrr6s zTriStageLRSchedule.__init__cCsp||jkr|j|j|}n||j|jkr|j}n|jt|j ||j|j}|jD]}||d<q/dS)zLCalculate the learning rate corresponding to the current step (num_updates).rN) rrrrrrurvrr )r r&rrrrrrr(Ms   zTriStageLRSchedule.__call__c Cs>|j|j|j|j|j|j|j|j|j|j d }t ||dS)r*) rrrrr!rrrrrN) rrrrr!rrrrrr+r,r-rrrr,_s zTriStageLRSchedule.saveFNcCsv~~t|}|d|_|d|_|d|_|d|_|d|_|d|_|d|_|d|_ |d |_ |d |_ d S) r1rrrrr!rrrrrN) r+r2rrrrr!rrrrrr3rrrr2ps          zTriStageLRSchedule.load)rrr6r7rrr"rrs; rrB)r;rur+rspeechbrain.utilsrspeechbrain.utils.loggerrr8r rregister_checkpoint_hooksrr?rXrdrlrrrrrrrrrModulerrrrrrsL     &Wa1ZN_{eaq+k1r^