How do I save a yolov5 model among the registered models in MLflow and then load it either directly or from weights (.pt) to make predictions?

Question

My objective is to :

1. save the yolov5 model in MLflow Model Registry
2. Load the model directly from model registry so that we can make predictions.
3. Have access to the results of prediction (bounding boxes).

Current steps followed:

1. Cloned the yolov5 repo (https://github.com/ultralytics/yolov5)
2. Modified the train.py file in yolov5 folder by including code that logs parameters to MLflow.
3. Executed mlflow ui only to find that the model has not got logged.

I followed the steps mentioned in this blog : https://techblog.sms-digital.com/integrating-mlflow-into-ultralytics/yolov5-1

Doing this properly created the experiment and logged the parameters and weights in the experiment section, but it did not log the model in the models section.

Metrics and parameters getting logged Models not getting logged

So I added the some lines of code in addition to the steps suggested in the blog. These lines were added in the main() function of train.py file of yolov5.

# Train
    if not opt.evolve:
        #---------------------------------------------------------------------------(added by me)
        with mlflow.start_run(run_name=opt.name, nested=True):
            _, model = train(opt.hyp, opt, device, callbacks) #changed the return values
        #---------------------------------------------------------------------------

    mlflow.pytorch.log_model(
        model,
        'yolov5_cloned_exp',
        registered_model_name='yolo_cloned_github'
    )

Even after doing this it did not work and only parameters were logged in the experiment section as before and Models were not registered.

Can you please help me by

1. Helping me understand what am I doing incorrectly
2. Suggesting any alternative methods to save and load yolov5 model in mlflow
3. Providing any reference links that might be helpful

Answer 1

The blog that you referenced mentions:

Step 4: When MLflow is started, log the model, if necessary. Note: scikit-learn, torch and keras models can be logged by log_model, but this was not possible with YOLOv5. Instead, we log the models as MLflow artifacts.