Source code for pytorch_lightning.callbacks.device_stats_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."""Device Stats Monitor====================Monitors and logs device stats during training."""fromtypingimportAny,Dict,Optionalimportpytorch_lightningasplfrompytorch_lightning.callbacks.baseimportCallbackfrompytorch_lightning.utilities.exceptionsimportMisconfigurationExceptionfrompytorch_lightning.utilities.typesimportSTEP_OUTPUTfrompytorch_lightning.utilities.warningsimportrank_zero_deprecation
[docs]classDeviceStatsMonitor(Callback):r""" Automatically monitors and logs device stats during training stage. ``DeviceStatsMonitor`` is a special callback as it requires a ``logger`` to passed as argument to the ``Trainer``. Raises: MisconfigurationException: If ``Trainer`` has no logger. Example: >>> from pytorch_lightning import Trainer >>> from pytorch_lightning.callbacks import DeviceStatsMonitor >>> device_stats = DeviceStatsMonitor() # doctest: +SKIP >>> trainer = Trainer(callbacks=[device_stats]) # doctest: +SKIP """
[docs]defsetup(self,trainer:"pl.Trainer",pl_module:"pl.LightningModule",stage:Optional[str]=None)->None:ifnottrainer.loggers:raiseMisconfigurationException("Cannot use DeviceStatsMonitor callback with Trainer that has no logger.")
[docs]defon_train_batch_start(self,trainer:"pl.Trainer",pl_module:"pl.LightningModule",batch:Any,batch_idx:int,unused:int=0,)->None:ifnottrainer.loggers:raiseMisconfigurationException("Cannot use `DeviceStatsMonitor` callback with `Trainer(logger=False)`.")ifnottrainer._logger_connector.should_update_logs:returndevice=trainer.strategy.root_devicedevice_stats=trainer.accelerator.get_device_stats(device)forloggerintrainer.loggers:separator=logger.group_separatorprefixed_device_stats=_prefix_metric_keys(device_stats,"on_train_batch_start",separator)logger.log_metrics(prefixed_device_stats,step=trainer.fit_loop.epoch_loop._batches_that_stepped)
[docs]defon_train_batch_end(self,trainer:"pl.Trainer",pl_module:"pl.LightningModule",outputs:STEP_OUTPUT,batch:Any,batch_idx:int,unused:int=0,)->None:ifnottrainer.loggers:raiseMisconfigurationException("Cannot use `DeviceStatsMonitor` callback with `Trainer(logger=False)`.")ifnottrainer._logger_connector.should_update_logs:returndevice=trainer.strategy.root_devicedevice_stats=trainer.accelerator.get_device_stats(device)forloggerintrainer.loggers:separator=logger.group_separatorprefixed_device_stats=_prefix_metric_keys(device_stats,"on_train_batch_end",separator)logger.log_metrics(prefixed_device_stats,step=trainer.fit_loop.epoch_loop._batches_that_stepped)
def_prefix_metric_keys(metrics_dict:Dict[str,float],prefix:str,separator:str)->Dict[str,float]:return{prefix+separator+k:vfork,vinmetrics_dict.items()}defprefix_metric_keys(metrics_dict:Dict[str,float],prefix:str)->Dict[str,float]:rank_zero_deprecation("`pytorch_lightning.callbacks.device_stats_monitor.prefix_metrics`"" is deprecated in v1.6 and will be removed in v1.8.")sep=""return_prefix_metric_keys(metrics_dict,prefix,sep)
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.