Skip to content

Register now for

# Reference for `ultralytics/utils/callbacks/mlflow.py`

Note

This file is available at https://github.com/ultralytics/ultralytics/blob/main/ultralytics/utils/callbacks/mlflow.py. If you spot a problem please help fix it by contributing a Pull Request 🛠️. Thank you 🙏!

## `` ultralytics.utils.callbacks.mlflow.sanitize_dict

    sanitize_dict(x: dict) -> dict

Sanitize dictionary keys by removing parentheses and converting values to floats.

Source code in `ultralytics/utils/callbacks/mlflow.py`

    42
    43
    44
    def sanitize_dict(x: dict) -> dict:
        """Sanitize dictionary keys by removing parentheses and converting values to floats."""
        return {k.replace("(", "").replace(")", ""): float(v) for k, v in x.items()}

## `` ultralytics.utils.callbacks.mlflow.on_pretrain_routine_end

    on_pretrain_routine_end(trainer)

Log training parameters to MLflow at the end of the pretraining routine.

This function sets up MLflow logging based on environment variables and trainer arguments. It sets the tracking URI, experiment name, and run name, then starts the MLflow run if not already active. It finally logs the parameters from the trainer.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`trainer` | `BaseTrainer` | The training object with arguments and parameters to log. | required

Environment Variables

MLFLOW_TRACKING_URI: The URI for MLflow tracking. If not set, defaults to 'runs/mlflow'. MLFLOW_EXPERIMENT_NAME: The name of the MLflow experiment. If not set, defaults to trainer.args.project. MLFLOW_RUN: The name of the MLflow run. If not set, defaults to trainer.args.name. MLFLOW_KEEP_RUN_ACTIVE: Boolean indicating whether to keep the MLflow run active after training ends.

Source code in `ultralytics/utils/callbacks/mlflow.py`

    47
    48
    49
    50
    51
    52
    53
    54
    55
    56
    57
    58
    59
    60
    61
    62
    63
    64
    65
    66
    67
    68
    69
    70
    71
    72
    73
    74
    75
    76
    77
    78
    79
    80
    81
    82
    83
    84
    85
    def on_pretrain_routine_end(trainer):
        """
        Log training parameters to MLflow at the end of the pretraining routine.

        This function sets up MLflow logging based on environment variables and trainer arguments. It sets the tracking URI,
        experiment name, and run name, then starts the MLflow run if not already active. It finally logs the parameters
        from the trainer.

        Args:
            trainer (ultralytics.engine.trainer.BaseTrainer): The training object with arguments and parameters to log.

        Environment Variables:
            MLFLOW_TRACKING_URI: The URI for MLflow tracking. If not set, defaults to 'runs/mlflow'.
            MLFLOW_EXPERIMENT_NAME: The name of the MLflow experiment. If not set, defaults to trainer.args.project.
            MLFLOW_RUN: The name of the MLflow run. If not set, defaults to trainer.args.name.
            MLFLOW_KEEP_RUN_ACTIVE: Boolean indicating whether to keep the MLflow run active after training ends.
        """
        global mlflow

        uri = os.environ.get("MLFLOW_TRACKING_URI") or str(RUNS_DIR / "mlflow")
        LOGGER.debug(f"{PREFIX} tracking uri: {uri}")
        mlflow.set_tracking_uri(uri)

        # Set experiment and run names
        experiment_name = os.environ.get("MLFLOW_EXPERIMENT_NAME") or trainer.args.project or "/Shared/Ultralytics"
        run_name = os.environ.get("MLFLOW_RUN") or trainer.args.name
        mlflow.set_experiment(experiment_name)

        mlflow.autolog()
        try:
            active_run = mlflow.active_run() or mlflow.start_run(run_name=run_name)
            LOGGER.info(f"{PREFIX}logging run_id({active_run.info.run_id}) to {uri}")
            if Path(uri).is_dir():
                LOGGER.info(f"{PREFIX}view at http://127.0.0.1:5000 with 'mlflow server --backend-store-uri {uri}'")
            LOGGER.info(f"{PREFIX}disable with 'yolo settings mlflow=False'")
            mlflow.log_params(dict(trainer.args))
        except Exception as e:
            LOGGER.warning(f"{PREFIX}Failed to initialize: {e}")
            LOGGER.warning(f"{PREFIX}Not tracking this run")

## `` ultralytics.utils.callbacks.mlflow.on_train_epoch_end

    on_train_epoch_end(trainer)

Log training metrics at the end of each train epoch to MLflow.

Source code in `ultralytics/utils/callbacks/mlflow.py`

    88
    89
    90
    91
    92
    93
    94
    95
    96
    97
    def on_train_epoch_end(trainer):
        """Log training metrics at the end of each train epoch to MLflow."""
        if mlflow:
            mlflow.log_metrics(
                metrics={
                    **sanitize_dict(trainer.lr),
                    **sanitize_dict(trainer.label_loss_items(trainer.tloss, prefix="train")),
                },
                step=trainer.epoch,
            )

## `` ultralytics.utils.callbacks.mlflow.on_fit_epoch_end

    on_fit_epoch_end(trainer)

Log training metrics at the end of each fit epoch to MLflow.

Source code in `ultralytics/utils/callbacks/mlflow.py`

    100
    101
    102
    103
    def on_fit_epoch_end(trainer):
        """Log training metrics at the end of each fit epoch to MLflow."""
        if mlflow:
            mlflow.log_metrics(metrics=sanitize_dict(trainer.metrics), step=trainer.epoch)

## `` ultralytics.utils.callbacks.mlflow.on_train_end

    on_train_end(trainer)

Log model artifacts at the end of training.

Source code in `ultralytics/utils/callbacks/mlflow.py`

    106
    107
    108
    109
    110
    111
    112
    113
    114
    115
    116
    117
    118
    119
    120
    121
    122
    123
    def on_train_end(trainer):
        """Log model artifacts at the end of training."""
        if not mlflow:
            return
        mlflow.log_artifact(str(trainer.best.parent))  # log save_dir/weights directory with best.pt and last.pt
        for f in trainer.save_dir.glob("*"):  # log all other files in save_dir
            if f.suffix in {".png", ".jpg", ".csv", ".pt", ".yaml"}:
                mlflow.log_artifact(str(f))
        keep_run_active = os.environ.get("MLFLOW_KEEP_RUN_ACTIVE", "False").lower() == "true"
        if keep_run_active:
            LOGGER.info(f"{PREFIX}mlflow run still alive, remember to close it using mlflow.end_run()")
        else:
            mlflow.end_run()
            LOGGER.debug(f"{PREFIX}mlflow run ended")

        LOGGER.info(
            f"{PREFIX}results logged to {mlflow.get_tracking_uri()}\n{PREFIX}disable with 'yolo settings mlflow=False'"
        )

📅 Created 1 year ago ✏️ Updated 1 year ago