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:
- Run Mlflow ui
pip install mlflow==1.26.1
mlflow ui
- 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:
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
🌟 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
MlflowLoggerclass 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.pymodule underutils/loggers/mlflowdirectory.
🎯 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.
📈 🔍 ⏱ ️
