geowatch.utils.lightning_ext.callbacks.telemetry module

LightningTelemetry callback to interface with torch.package

class geowatch.utils.lightning_ext.callbacks.telemetry.LightningTelemetry

Bases: Callback

The idea is that we wrap a fit job with ProcessContext

Example

>>> from geowatch.utils.lightning_ext.callbacks.telemetry import *  # NOQA
>>> from geowatch.utils.lightning_ext.demo import LightningToyNet2d
>>> from geowatch.utils.lightning_ext.callbacks import StateLogger
>>> import ubelt as ub
>>> from geowatch.utils.lightning_ext.monkeypatches import disable_lightning_hardware_warnings
>>> disable_lightning_hardware_warnings()
>>> default_root_dir = ub.Path.appdir('geowatch/lightning_ext/test/telemetry')
>>> default_root_dir.delete().ensuredir()
>>> self = LightningTelemetry()
>>> # Test starting a model without any existing checkpoints
>>> trainer = pl.Trainer(default_root_dir=default_root_dir, callbacks=[
>>>     self,
>>>     StateLogger()
>>> ], max_epochs=2, accelerator='cpu', devices=1)
>>> model = LightningToyNet2d()
>>> trainer.fit(model)

classmethod add_argparse_args(parent_parser)

Example

>>> from geowatch.utils.lightning_ext.callbacks.telemetry import *  # NOQA
>>> from geowatch.utils.configargparse_ext import ArgumentParser
>>> cls = LightningTelemetry
>>> parent_parser = ArgumentParser(formatter_class='defaults')
>>> cls.add_argparse_args(parent_parser)
>>> parent_parser.print_help()

setup(trainer, pl_module, stage=None)

Finalize initialization step. Resolve the paths where files will be written.

Parameters:

• trainer (pl.Trainer)

• pl_module (pl.LightningModule)

• stage (str | None)

Returns:

None

on_fit_start(trainer: Trainer, pl_module: LightningModule) → None

Todo

• [ ] Write out the uninitialized topology

on_fit_end(trainer: Trainer, pl_module: LightningModule) → None

on_train_epoch_end(trainer, logs=None)

on_exception(trainer: Trainer, pl_module: LightningModule, *args, **kw) → None