Source code for pytorch_lightning.callbacks.lr_monitor
# Copyright The PyTorch Lightning team.## Licensed under the Apache License, Version 2.0 (the "License");# you may not use this file except in compliance with the License.# You may obtain a copy of the License at## http://www.apache.org/licenses/LICENSE-2.0## Unless required by applicable law or agreed to in writing, software# distributed under the License is distributed on an "AS IS" BASIS,# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.# See the License for the specific language governing permissions and# limitations under the License.r"""Learning Rate Monitor=====================Monitor and logs learning rate for lr schedulers during training."""importitertoolsfromcollectionsimportdefaultdictfromtypingimportAny,DefaultDict,Dict,List,Optional,Set,Tuple,Typefromtorch.optim.optimizerimportOptimizerimportpytorch_lightningasplfrompytorch_lightning.callbacks.baseimportCallbackfrompytorch_lightning.utilities.exceptionsimportMisconfigurationExceptionfrompytorch_lightning.utilities.rank_zeroimportrank_zero_deprecation,rank_zero_warnfrompytorch_lightning.utilities.typesimportLRSchedulerConfig
[docs]classLearningRateMonitor(Callback):r""" Automatically monitor and logs learning rate for learning rate schedulers during training. Args: logging_interval: set to ``'epoch'`` or ``'step'`` to log ``lr`` of all optimizers at the same interval, set to ``None`` to log at individual interval according to the ``interval`` key of each scheduler. Defaults to ``None``. log_momentum: option to also log the momentum values of the optimizer, if the optimizer has the ``momentum`` or ``betas`` attribute. Defaults to ``False``. Raises: MisconfigurationException: If ``logging_interval`` is none of ``"step"``, ``"epoch"``, or ``None``. Example:: >>> from pytorch_lightning import Trainer >>> from pytorch_lightning.callbacks import LearningRateMonitor >>> lr_monitor = LearningRateMonitor(logging_interval='step') >>> trainer = Trainer(callbacks=[lr_monitor]) Logging names are automatically determined based on optimizer class name. In case of multiple optimizers of same type, they will be named ``Adam``, ``Adam-1`` etc. If a optimizer has multiple parameter groups they will be named ``Adam/pg1``, ``Adam/pg2`` etc. To control naming, pass in a ``name`` keyword in the construction of the learning rate schedulers. A ``name`` keyword can also be used for parameter groups in the construction of the optimizer. Example:: def configure_optimizer(self): optimizer = torch.optim.Adam(...) lr_scheduler = { 'scheduler': torch.optim.lr_scheduler.LambdaLR(optimizer, ...) 'name': 'my_logging_name' } return [optimizer], [lr_scheduler] Example:: def configure_optimizer(self): optimizer = torch.optim.SGD( [{ 'params': [p for p in self.parameters()], 'name': 'my_parameter_group_name' }], lr=0.1 ) lr_scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, ...) return [optimizer], [lr_scheduler] """def__init__(self,logging_interval:Optional[str]=None,log_momentum:bool=False)->None:iflogging_intervalnotin(None,"step","epoch"):raiseMisconfigurationException("logging_interval should be `step` or `epoch` or `None`.")self.logging_interval=logging_intervalself.log_momentum=log_momentumself.lrs:Dict[str,List[float]]={}self._lr_sch_names:List[str]=[]
[docs]defon_train_start(self,trainer:"pl.Trainer",*args:Any,**kwargs:Any)->None:"""Called before training, determines unique names for all lr schedulers in the case of multiple of the same type or in the case of multiple parameter groups. Raises: MisconfigurationException: If ``Trainer`` has no ``logger``. """ifnottrainer.loggers:raiseMisconfigurationException("Cannot use `LearningRateMonitor` callback with `Trainer` that has no logger.")ifself.log_momentum:def_check_no_key(key:str)->bool:iftrainer.lr_scheduler_configs:returnany(keynotinconfig.scheduler.optimizer.defaultsforconfigintrainer.lr_scheduler_configs)returnany(keynotinoptimizer.defaultsforoptimizerintrainer.optimizers)if_check_no_key("momentum")and_check_no_key("betas"):rank_zero_warn("You have set log_momentum=True, but some optimizers do not"" have momentum. This will log a value 0 for the momentum.",category=RuntimeWarning,)# Find names for schedulersnames:List[List[str]]=[](sched_hparam_keys,optimizers_with_scheduler,optimizers_with_scheduler_types,)=self._find_names_from_schedulers(trainer.lr_scheduler_configs)names.extend(sched_hparam_keys)# Find names for leftover optimizersoptimizer_hparam_keys,_=self._find_names_from_optimizers(trainer.optimizers,seen_optimizers=optimizers_with_scheduler,seen_optimizer_types=optimizers_with_scheduler_types,)names.extend(optimizer_hparam_keys)# Initialize for storing valuesnames_flatten=list(itertools.chain.from_iterable(names))self.lrs={name:[]fornameinnames_flatten}self.last_momentum_values={name+"-momentum":Nonefornameinnames_flatten}
def_extract_stats(self,trainer:"pl.Trainer",interval:str)->Dict[str,float]:latest_stat={}(scheduler_hparam_keys,optimizers_with_scheduler,optimizers_with_scheduler_types,)=self._find_names_from_schedulers(trainer.lr_scheduler_configs,add_lr_sch_names=False)self._remap_keys(scheduler_hparam_keys)forname,configinzip(scheduler_hparam_keys,trainer.lr_scheduler_configs):ifintervalin[config.interval,"any"]:opt=config.scheduler.optimizercurrent_stat=self._get_lr_momentum_stat(opt,name)latest_stat.update(current_stat)optimizer_hparam_keys,optimizers_without_scheduler=self._find_names_from_optimizers(trainer.optimizers,seen_optimizers=optimizers_with_scheduler,seen_optimizer_types=optimizers_with_scheduler_types,add_lr_sch_names=False,)self._remap_keys(optimizer_hparam_keys)foropt,namesinzip(optimizers_without_scheduler,optimizer_hparam_keys):current_stat=self._get_lr_momentum_stat(opt,names)latest_stat.update(current_stat)returnlatest_statdef_get_lr_momentum_stat(self,optimizer:Optimizer,names:List[str])->Dict[str,float]:lr_momentum_stat={}param_groups=optimizer.param_groupsuse_betas="betas"inoptimizer.defaultsforpg,nameinzip(param_groups,names):lr=self._extract_lr(pg,name)lr_momentum_stat.update(lr)momentum=self._extract_momentum(param_group=pg,name=name.replace(name,f"{name}-momentum"),use_betas=use_betas)lr_momentum_stat.update(momentum)returnlr_momentum_statdef_extract_lr(self,param_group:Dict[str,Any],name:str)->Dict[str,Any]:lr=param_group["lr"]self.lrs[name].append(lr)return{name:lr}def_remap_keys(self,names:List[List[str]],token:str="/pg1")->None:"""This function is used the remap the keys if param groups for a given optimizer increased."""forgroup_new_namesinnames:fornew_nameingroup_new_names:old_name=new_name.replace(token,"")iftokeninnew_nameandold_nameinself.lrs:self.lrs[new_name]=self.lrs.pop(old_name)elifnew_namenotinself.lrs:self.lrs[new_name]=[]def_extract_momentum(self,param_group:Dict[str,List],name:str,use_betas:bool)->Dict[str,float]:ifnotself.log_momentum:return{}momentum=param_group["betas"][0]ifuse_betaselseparam_group.get("momentum",0)self.last_momentum_values[name]=momentumreturn{name:momentum}def_add_prefix(self,name:str,optimizer_cls:Type[Optimizer],seen_optimizer_types:DefaultDict[Type[Optimizer],int])->str:ifoptimizer_clsnotinseen_optimizer_types:returnnamecount=seen_optimizer_types[optimizer_cls]returnname+f"-{count-1}"ifcount>1elsenamedef_add_suffix(self,name:str,param_groups:List[Dict],param_group_index:int,use_names:bool=True)->str:iflen(param_groups)>1:ifnotuse_names:returnf"{name}/pg{param_group_index+1}"pg_name=param_groups[param_group_index].get("name",f"pg{param_group_index+1}")returnf"{name}/{pg_name}"elifuse_names:pg_name=param_groups[param_group_index].get("name")returnf"{name}/{pg_name}"ifpg_nameelsenamereturnnamedef_duplicate_param_group_names(self,param_groups:List[Dict])->Set[str]:names=[pg.get("name",f"pg{i}")fori,pginenumerate(param_groups,start=1)]unique=set(names)iflen(names)==len(unique):returnset()return{nforninnamesifnames.count(n)>1}def_find_names_from_schedulers(self,lr_scheduler_configs:List[LRSchedulerConfig],add_lr_sch_names:bool=True)->Tuple[List[List[str]],List[Optimizer],DefaultDict[Type[Optimizer],int]]:# Create unique names in the case we have multiple of the same learning# rate scheduler + multiple parameter groupsnames=[]seen_optimizers:List[Optimizer]=[]seen_optimizer_types:DefaultDict[Type[Optimizer],int]=defaultdict(int)forconfiginlr_scheduler_configs:sch=config.schedulerifconfig.nameisnotNone:name=config.nameelse:name="lr-"+sch.optimizer.__class__.__name__updated_names=self._check_duplicates_and_update_name(sch.optimizer,name,seen_optimizers,seen_optimizer_types,config,add_lr_sch_names)names.append(updated_names)returnnames,seen_optimizers,seen_optimizer_typesdef_find_names_from_optimizers(self,optimizers:List[Any],seen_optimizers:List[Optimizer],seen_optimizer_types:DefaultDict[Type[Optimizer],int],add_lr_sch_names:bool=True,)->Tuple[List[List[str]],List[Optimizer]]:names=[]optimizers_without_scheduler=[]foroptimizerinoptimizers:# Deepspeed optimizer wraps the native optimizeroptimizer=optimizer.optimizerifhasattr(optimizer,"optimizer")elseoptimizerifoptimizerinseen_optimizers:continuename="lr-"+optimizer.__class__.__name__updated_names=self._check_duplicates_and_update_name(optimizer,name,seen_optimizers,seen_optimizer_types,None,add_lr_sch_names)names.append(updated_names)optimizers_without_scheduler.append(optimizer)returnnames,optimizers_without_schedulerdef_check_duplicates_and_update_name(self,optimizer:Optimizer,name:str,seen_optimizers:List[Optimizer],seen_optimizer_types:DefaultDict[Type[Optimizer],int],lr_scheduler_config:Optional[LRSchedulerConfig],add_lr_sch_names:bool=True,)->List[str]:seen_optimizers.append(optimizer)optimizer_cls=type(optimizer)iflr_scheduler_configisnotNoneandlr_scheduler_config.nameisNone:seen_optimizer_types[optimizer_cls]+=1eliflr_scheduler_configisNone:seen_optimizer_types[optimizer_cls]+=1# Multiple param groups for the same optimizerparam_groups=optimizer.param_groupsduplicates=self._duplicate_param_group_names(param_groups)ifduplicates:raiseMisconfigurationException("A single `Optimizer` cannot have multiple parameter groups with identical "f"`name` values. {name} has duplicated parameter group names {duplicates}")name=self._add_prefix(name,optimizer_cls,seen_optimizer_types)name_list=[self._add_suffix(name,param_groups,i)foriinrange(len(param_groups))]ifadd_lr_sch_names:self._lr_sch_names.append(name)returnname_list@propertydeflr_sch_names(self)->List[str]:# TODO remove `lr_sch_names` and `add_lr_sch_names` argument in v1.7.0rank_zero_deprecation("`LearningRateMonitor.lr_sch_names` has been deprecated in v1.5 and will be removed in 1.7."" Consider accessing them using `LearningRateMonitor.lrs.keys()` which will return"" the names of all the optimizers, even those without a scheduler.")returnself._lr_sch_names
To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. Read PyTorch Lightning's Privacy Policy.