neural_lam.custom_loggers

neural_lam.custom_loggers#

Custom logging utilities (e.g., MLFlow wrappers) used in Neural-LAM.

Module Contents#

class neural_lam.custom_loggers.CustomMLFlowLogger(experiment_name, tracking_uri, run_name)#

Bases: pytorch_lightning.loggers.MLFlowLogger

Custom MLFlow logger that adds the log_image() functionality not present in the default implementation from pytorch-lightning as of version 2.0.3 at least.

Initialize the logger and start an MLflow run.

Parameters:

  • experiment_name (str) – Target MLflow experiment.

  • tracking_uri (str) – MLflow tracking server URI.

  • run_name (str) – Human-readable run name stored as mlflow.runName.

log_image(key, images, step=None)#

Log one or more Matplotlib figures as images in MLflow.

Parameters:

  • key (str) – Identifier under which to log the image.

  • images (Sequence[matplotlib.figure.Figure]) – Figures to export; only the first element is logged.

  • step (int or None, optional) – Optional training step index appended to key.

Raises:

SystemExit – If AWS credentials for the MLflow artifact store are missing.

property save_dir#

Returns the directory where the MLFlow artifacts are saved. Used to define the path to save output when using the logger.

Returns:

Path to the directory where the artifacts are saved.

Return type:

str

Contents