OpenBot
OpenBot copied to clipboard
Add a pipeline to train yolov5 object-detection model on custom data
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
-
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...)
-
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)
This looks great @mnrozhkov. Looking forward to a PR in the future.