MLFlowLogger

class lightning.pytorch.loggers.MLFlowLogger(experiment_name='lightning_logs', run_name=None, tracking_uri=None, tags=None, save_dir='./mlruns', log_model=False, prefix='', artifact_location=None, run_id=None, synchronous=None)

Bases: Logger

Log using MLflow.

Install it with pip:

pip install mlflow  # or mlflow-skinny

from lightning.pytorch import Trainer
from lightning.pytorch.loggers import MLFlowLogger

mlf_logger = MLFlowLogger(experiment_name="lightning_logs", tracking_uri="file:./ml-runs")
trainer = Trainer(logger=mlf_logger)

Use the logger anywhere in your LightningModule as follows:

from lightning.pytorch import LightningModule


class LitModel(LightningModule):
    def training_step(self, batch, batch_idx):
        # example
        self.logger.experiment.whatever_ml_flow_supports(...)

    def any_lightning_module_function_or_hook(self):
        self.logger.experiment.whatever_ml_flow_supports(...)

Parameters:

• experiment_name (str) – The name of the experiment.

• run_name (Optional[str]) – Name of the new run. The run_name is internally stored as a mlflow.runName tag. If the mlflow.runName tag has already been set in tags, the value is overridden by the run_name.

• tracking_uri (Optional[str]) – Address of local or remote tracking server. If not provided, defaults to MLFLOW_TRACKING_URI environment variable if set, otherwise it falls back to file:<save_dir>.

• tags (Optional[dict[str, Any]]) – A dictionary tags for the experiment.

• save_dir (Optional[str]) – A path to a local directory where the MLflow runs get saved. Defaults to ./mlruns if tracking_uri is not provided. Has no effect if tracking_uri is provided.

• log_model (Literal[True, False, 'all']) –

  Log checkpoints created by ModelCheckpoint as MLFlow artifacts.

  • if log_model == 'all', checkpoints are logged during training.

  • if log_model == True, checkpoints are logged at the end of training, except when save_top_k == -1 which also logs every checkpoint during training.

  • if log_model == False (default), no checkpoint is logged.

• prefix (str) – A string to put at the beginning of metric keys.

• artifact_location (Optional[str]) – The location to store run artifacts. If not provided, the server picks an appropriate default.

• run_id (Optional[str]) – The run identifier of the experiment. If not provided, a new run is started.

• synchronous (Optional[bool]) – Hints mlflow whether to block the execution for every logging call until complete where applicable. Requires mlflow >= 2.8.0

Raises:

ModuleNotFoundError – If required MLFlow package is not installed on the device.

Note

MLFlowLogger will skip logging any metric (same name and step) more than once per run, to prevent database unique constraint violations on some MLflow backends (such as PostgreSQL). Only the first value for each (metric, step) pair will be logged per run.

after_save_checkpoint(checkpoint_callback)

Called after model checkpoint callback saves a new checkpoint.

Parameters:

checkpoint_callback (ModelCheckpoint) – the model checkpoint callback instance

Return type:

None

finalize(status='success')

Do any processing that is necessary to finalize an experiment.

Parameters:

status (str) – Status that the experiment finished with (e.g. success, failed, aborted)

Return type:

None

log_hyperparams(params)

Record hyperparameters.

Parameters:

• params (Union[dict[str, Any], Namespace]) – Namespace or Dict containing the hyperparameters

• args – Optional positional arguments, depends on the specific logger being used

• kwargs – Optional keyword arguments, depends on the specific logger being used

Return type:

None

log_metrics(metrics, step=None)

Records metrics. This method logs metrics as soon as it received them.

Parameters:

• metrics (Mapping[str, float]) – Dictionary with metric names as keys and measured quantities as values

• step (Optional[int]) – Step number at which the metrics should be recorded

Return type:

None

property experiment: mlflow.tracking.MlflowClient

Actual MLflow object. To use MLflow features in your LightningModule do the following.

Example:

self.logger.experiment.some_mlflow_function()

property experiment_id: Optional[str]

Create the experiment if it does not exist to get the experiment id.

Returns:

The experiment id.

property name: Optional[str]

Get the experiment id.

Returns:

The experiment id.

property run_id: Optional[str]

Create the experiment if it does not exist to get the run id.

Returns:

The run id.

property save_dir: Optional[str]

The root file directory in which MLflow experiments are saved.

Returns:

Local path to the root experiment directory if the tracking uri is local. Otherwise returns None.

property version: Optional[str]

Get the run id.

Returns:

The run id.