OpenBot icon indicating copy to clipboard operation
OpenBot copied to clipboard

Add a pipeline to train yolov5 object-detection model on custom data

Open mnrozhkov opened this issue 2 years ago • 1 comments

Is your feature request related to a problem? Please describe. Sometimes users may want to train a custom object-detection model with YoloV5 and use it in the OpenBot. To make it simple, we may provide a pipeline to collect/annotate/train/deploy model to the robot. This request is focusing of the model training pipeline.

Describe the solution you'd like Proposed solution utilizes DVC to run end-to-end YoloV5 model training in Docker

  • DVC pipeline (`dvc.yaml) has three stages: download_model, train, val and exports
    • download_model: download pre-train model
    • train: train a model
    • val: run validation checks
    • exports: export models into required formats
  • the pipeline configuration is in params.yaml file

DVC allows to automate model training & validation. With DVC remote storage setup, user may store models in local or cloud storages. Using DVC users may switch between different experiments and model versions

Prototype: https://github.com/mnrozhkov/OpenBot/tree/dev-object-detection (in progress)

Describe alternatives you've considered

  1. There are two alternatives on YoloV5 dev environment setup: using virtual environment or Docker. Docker looks more appealing because it make environment setup simple and reproducible. Also, it's easier to control location of output artefacts (models, plots...)

  2. The pipeline can be moved into a separate repository. It may be easier to experiment and share models.

Additional context

Next steps/ideas:

  • add a pipeline to prepare/link custom dataset (share dataset?)
  • add CI configuration to run model training in Cloud
  • add Model Registry (to make it easy to plug and play with different models)

mnrozhkov avatar Nov 15 '22 19:11 mnrozhkov

This looks great @mnrozhkov. Looking forward to a PR in the future.

thias15 avatar Nov 16 '22 13:11 thias15