feat: Adds mlflow run logging

Created by: ElefHead

Saw https://github.com/ultralytics/yolov5/issues/4231 and https://github.com/ultralytics/yolov5/pull/7840 and figured I'd make a PR to add some rudimentary logging for Mlflow that I've successfully used.

  • Assumes that mlflow and experiment is setup by user
  • Logger creates a run and tracks params, metrics, and artifacts

To make use of this code:

  1. Run Mlflow ui
pip install mlflow==1.26.1
mlflow ui
  1. Run training code. eg.
python train.py --img 640 --batch 16 --epochs 3 --data coco128.yaml --weights yolov5n.pt --cache cache --save-period 1

This will create a run and log all params, metrics, artifacts, and log the model to the run. If no experiment is set, everything will be logged to the default experiment. Below is an example of a run: Screen Shot 2022-08-25 at 11 09 26 PM

Alternatively, users can also run the code as an Mlflow Project. For that users would have to create their own MLproject file, specify the backend-config and conda.yaml and run the project as

mlflow run yolov5 -p {param1} --experiment-name {experiment-name} -b {databricks or local} --backend-config {set only if running on databricks}

and all of this will still work.

🛠️ PR Summary

Made with ️ by Ultralytics Actions

🌟 Summary

Integration of Mlflow for experiment tracking in Ultralytics YOLOv5.

📊 Key Changes

  • Added Mlflow (mlflow==1.26.1) to requirements, expanding the suite of supported experiment tracking tools.
  • Introduced MlflowLogger class to handle Mlflow logging activities within the code.
  • Enabled artifact logging (like labels, mosaics, and validation images), parameters, metrics, and model logging through Mlflow.
  • Updated logger initialization code to support the new Mlflow integration.
  • Created dedicated mlflow_utils.py module under utils/loggers/mlflow directory.

🎯 Purpose & Impact

  • Purpose: To provide users additional flexibility to track and visualize their training experiments using Mlflow.
  • Impact: Users can now choose Mlflow as their experiment tracking tool, potentially leveraging existing Mlflow setups and integrating with Mlflow's extensive ecosystem. This enhances the experiment management and reproducibility features of YOLOv5 training workflows. 📈🔍

Merge request reports

Loading