Train SOTA Object Detection Models With DagsHub and Deci AI Super Gradients

DagsHub is now integrated with Deci to enable experiment logging with MLflow and DVC when training SOTA object detection models. Achieve full reproducibility and model <-> data lineage with a few lines of code.

We see great value in the recent development by Deci, not only in improving object detection capabilities but also in showcasing the deep learning models that can improve other models using technologies like neural architecture search (NAS) and Automated Neural Architecture Construction (AutoNAC).

That’s why we decided to extend its capabilities and support out-of-the-box logging using a DagsHub Logger callback for the new YOLO-NAS model. With this integration, you can log your experiments and artifacts to DagsHub remote servers, using MLflow and DVC under the hood, with minor changes to your code.

Want to learn how you can use it? This blog is for you!

What is YOLO-NAS?

YOLO-NAS is a new object detection model developed by researchers at Deci, using deep learning models. If you want to dive into its implementation details and performance, we highly recommend reading the launch blog.

What does the integration between DagsHub and Deci include?

With the new integration between DagsHub and Deci, you can now log experiments and trained models to your remote MLflow server hosted on DagsHub, diff experiments, and share them with your friends and colleagues. If you’re using DVC, this integration enables you to version the trained model with DVC and push it to your DagsHub remote storage. All these are encapsulated under the DagsHub Logger and only require adding a few lines of code to your project to unlock a world of new capabilities.

How to set up DagsHub Logger with Deci?

DagsHub logger is supported as a callback in SuperGradients and can be used by adding a few parameters to the trainer module.

Set DagsHub as the trainer's logger:

training_params["sg_logger"] = "dagshub_sg_logger"

Provide the DagsHub repo name and owner:

:param dagshub_repository: Format: <dagshub_username>/<dagshub_reponame>

In the event that the repository does not exist, it will be created automatically on your behalf.

On top of that, you can also set the logger to version the trained model using DVC by setting the log_mlflow_only parameter to False. If you’d like to use it as part of an automated pipeline, you can set the dagshub_auth which resolves the authentication process. You can find further information about the logger parameters in the source code.

Comparison of YOLO-NAS vs. YOLOv5 using DagsHub Logger

We wanted to take YOLO-NAS for a test drive and compare its performance to YOLOv5. For that, we used the PlaneDetector project where we tried to detect aircraft in satellite images. This is a perfect project for this task as the objects in the images are fairly small which requires high accuracy of the models.

To test the models’ performances, we’ll use the DagsHub Logger to log the metrics and parameters during training and evaluation and perform a quality comparison between the two architectures using DagsHub’s Experiment table.

Train YOLO-NAS with DagsHub Logger

One of the great advantages of YOLO-NAS is that it make it easy to train from scratch or fine-tune new models and require writing very little code for that. In our project, we used the Trainer module from the SuperGradients Python package, and the build in data loders.

• Data loader configurations: our annotations are in COCO format, which is supported as part of the package and only require to provide the path to the files in the local file system.
• DagsHub logger configurations: We’ll start by defining the sg_logger as dagshub_sg_logger and provide the path to the DagsHub repo as DagsHub/PlaneDetector. Next, we want to log the artifacts with DVC, so we’ll set log_mlflow_only to False. Here are the full logger configurations:

"sg_logger": "dagshub_sg_logger",
    "sg_logger_params":  # Params that will be passes to __init__ of the logger super_gradients.common.sg_loggers.dagshub_sg_logger.DagsHubSGLogger
                     {
                       "dagshub_repository": "DagsHub/PlaneDetector", # Optional: Your DagsHub project name, consisting of the owner name, followed by '/', and the repo name. If this is left empty, you'll be prompted in your run to fill it in manually.
                       "log_mlflow_only": False, # Optional: Change to true to bypass logging to DVC, and log all artifacts only to MLflow  
                       "save_checkpoints_remote": True,
                       "save_tensorboard_remote": True,
                       "save_logs_remote": True,
                     },

You can find the full source code in the DagsHub repo.

We trained YOLO-NAS for 100 epochs and reached 0.948 mAP50, 0.992 Recall, and 0.954 precision. You can find the full logs of the training process under the experiment tab

Train YOLOv5 with DagsHub Logger

We utilized the training and validation code provided by Ultralytics for training the YOLOv5 detection model. To train with your customized datasets, please organize the data folder structure as Ultralytics suggested.  Now, we can easily start by training the YOLOv5 with the customized dataset using the command below:

python train_plane_detector.py --data data/airplane.yaml --weights ../models/yolov5s.pt --save-path ../models/ --epochs=100

The training command line supports the customized settings for the data YAML file, pre-trained model,  folder to save the trained models, and the maximum epochs for training.

Be noted that if you're using Colab Notebook for training, please specify the CUBLAS_WORKSPACE_CONFIG environment variables to prevent the error:

import os
os.environ["CUBLAS_WORKSPACE_CONFIG"]=":4096:8"

We trained YOLOv5 for 100 epochs and reached 0.946 mAP50, 0.898 Recall, and 0.953 precision. You can also find the full logs of the training process under the experiment tab

Result Comparison

We can see that YOLO-NAS provides (a slight) improvement on all fronts compared to YOLOv5. Get deeper into different metrics, both models have almost the same precision, which means they both make few false positive predictions. However, YOLO-NAS (0.992) has a much higher recall rate comparing to YOLO-V5 (0.898), which means it can detect more positive instances. The mAP@0.5 taking into account both precision and recall. YOLOv5 has a mAP@0.5 of 0.946 while YOLO-NAS has a mAP@0.5 of 0.948, indicating that the YOLO-NAS has a better overall performance comparing to YOLOv5.

Conclusions

The new integration between DagsHub and YOLO-NAS makes it easy for you to log experiments, artifacts, and collaborate with others on machine learning projects. Give the DagsHub Logger a try and see how it can enhance your machine learning workflow.

Let us know how it goes on our community Discord and if you have any enhancements requests - we’d love to enrich this integration and add more capabilities!