gnn_tracking.utils.lightning#

Attributes#

logger

Classes#

StandardError

A torch metric that computes the standard error.

SimpleTqdmProgressBar

Fallback progress bar that creates a new tqdm bar for each epoch.

Functions#

save_sub_hyperparameters(→ None)

Take hyperparameters from obj and save them to self under the

load_obj_from_hparams(→ Any)

Load object from hyperparameters.

obj_from_or_to_hparams(→ Any)

Used to support initializing python objects from hyperparameters:

get_object_from_path(→ Any)

Get object from path (string) to its code location.

get_lightning_module(...)

Get model (specified by class_path, a string) and

get_model(→ torch.nn.Module | None)

Get torch model (specified by class_path, a string) and load a checkpoint.

find_latest_checkpoint(→ pathlib.Path)

Find latest lightning checkpoint

Module Contents#

gnn_tracking.utils.lightning.logger#
gnn_tracking.utils.lightning.save_sub_hyperparameters(self: pytorch_lightning.core.mixins.hparams_mixin.HyperparametersMixin, key: str, obj: pytorch_lightning.core.mixins.hparams_mixin.HyperparametersMixin | dict, errors: Literal['warn', 'raise'] = 'warn') None#

Take hyperparameters from obj and save them to self under the key key.

Parameters:
  • self – The object to save the hyperparameters to.

  • key – The key under which to save the hyperparameters.

  • obj – The object to take the hyperparameters from.

  • errors – Whether to raise an error or just warn

gnn_tracking.utils.lightning.load_obj_from_hparams(hparams: dict[str, Any], key: str = '') Any#

Load object from hyperparameters.

gnn_tracking.utils.lightning.obj_from_or_to_hparams(self: pytorch_lightning.core.mixins.hparams_mixin.HyperparametersMixin, key: str, obj: Any) Any#

Used to support initializing python objects from hyperparameters: If obj is a python object other than a dictionary, its hyperparameters are saved (its class path and init args) to self.hparams[key]. If obj is instead a dictionary, its assumed that we have to restore an object based on this information.

gnn_tracking.utils.lightning.get_object_from_path(path: str, init_args: dict[str, Any] | None = None) Any#

Get object from path (string) to its code location.

gnn_tracking.utils.lightning.get_lightning_module(class_path: str, chkpt_path: str = '', *, freeze: bool = True, device: str | None = None) pytorch_lightning.LightningModule | None#

Get model (specified by class_path, a string) and load a checkpoint.

gnn_tracking.utils.lightning.get_model(class_path: str, chkpt_path: str = '', freeze: bool = False, whole_module: bool = False, device: None | str = None) torch.nn.Module | None#

Get torch model (specified by class_path, a string) and load a checkpoint. Uses get_lightning_module to get the model.

Parameters:
  • class_path – The path to the lightning module class

  • chkpt_path – The path to the checkpoint. If no checkpoint is specified, we return None.

  • freeze – Whether to freeze the model

  • whole_module – Whether to return the whole lightning module or just the model

  • device

class gnn_tracking.utils.lightning.StandardError#

Bases: torchmetrics.Metric

A torch metric that computes the standard error. This is necessary, because LightningModule.log doesn’t take custom reduce functions.

update(x: torch.Tensor)#
compute()#
class gnn_tracking.utils.lightning.SimpleTqdmProgressBar#

Bases: pytorch_lightning.callbacks.ProgressBar

Fallback progress bar that creates a new tqdm bar for each epoch. Adapted from Lightning-AI/lightning#2189 , reply Lightning-AI/lightning#2189

property is_enabled#
on_train_epoch_start(trainer, pl_module)#
on_train_batch_end(trainer, pl_module, outputs, batch, batch_idx)#
on_validation_epoch_end(trainer, pl_module) None#
disable()#
gnn_tracking.utils.lightning.find_latest_checkpoint(log_dir: os.PathLike, trial_name: str = '') pathlib.Path#

Find latest lightning checkpoint

Parameters:
  • log_dir (os.PathLike) – Path to the directory of your trial or to the directory of the experiment if trial_name is specified.

  • trial_name (str, optional) – Name of the trial if log_dir is the directory of the experiment.