The scripts in this repo can be used to export a YOLOv5 model to coreml and benchmark it.
- python 3.8.x
- poetry
Other dependencies are installed by poetry automatically with:
$ poetry install
Note: It assumes that you've cloned the yolov5 repo to ../yolov5
relative to the project and that you have added a setup.py
file, so poetry installs the dependencies of YOLOv5. See Issue #2525.
It's recommended to use the version as specified. CoreML Tools only works with certain PyTorch Versions, which are only available for certain Python versions. You might want to consider using pyenv:
$ pyenv install 3.8.6
$ pyenv global 3.8.6
$ poetry install
We use the unified conversion api of coremltools 4.0 to convert a YOLOv5 model to CoreML. Additionally we add an export layer, so the model integrates well with Apple's Vision Framework.
For optimal integration in Swift you need at least iOS 13
as the export layer with NMS is otherwise not supported. It is possible though to use iOS 12
if you implement NMS manually in Swift.
Experience has shown that one needs to be very careful with the versions of libaries used. In particular the PyTorch version. It's recommended to not change the version of the libaries used. If you need for some reasons a newer PyTorch version, check the github page of coremltools for open issues regarding the PyTorch version you want to use, in the past most recents versions where not compatible.
The models always need the original source code, unless you do have a torchscript model and therefore skip the tracing step in the script.
Note: It has a huge impact on performance if the model runs on the NeuralEngine or the CPU / GPU (or switches between them) on your device. Unfortunately, there is no documentation which model layers can run on the neural engine and which not (some infos here). With yolov5 version 2, 3 and 4 there were problems with the SPP Layers with kernel sizes bigger than 7, so we replaced them and retrained the model. On a recent device YOLOv5s should be around 20ms / detection. See Issue 2526
First, some values in the script should be changed according to your needs. In src/coreml_export/main.py
you'll find some global variables, which are specific to the concrete model you use:
classLabels
-> This is the list of labels your model recognized, all pretrained models of YOLOv5 used the coco dataset and therefore regonize 80 Labelsanchors
-> These depend on the model version you use (s, m, l, x) and you will find these in the accordingyolo<Version>.yml
file in theyolov5
repository (Use the files from the correct yolov5 version!).reverseModel
-> Some models have their strides and anchors in reversed order. This variable exists for convinience to quickly change the order of those.
To run the script use the command:
$ poetry run coreml-export
Run it with -h
to get a list of optional arguments to customize model input & output paths / names and other things.
There is a simple script to test the exported model with some sample images in src/coreml_export/test.py
. You should check if the predictions are similar to those of the original PyTorch model. Be aware that the predictions will be slightly different though. If there are huge differences this might be a hint that you need to set reverseModel
accordingly.
To run the script use the command:
$ poetry run coreml-test
Most important is that the model runs fully on the Neural Engine. There is no official documentation, but take a look at Everything we actually know about the Apple Neural Engine (ANE).
In src/coreml_export/snippets.py
you might find a few helpful snippets to temporarily (!) change layers, parameters or other things of the model and test how it influences performance.
This makes heavy use of the library Object Detection Metrics Library developed by @rafaelpadilla under the MIT License.
The library is included in the objectionDetectionMetrics
subfolder with some small adjustments.
The Metrics script in src/coreml_metrics/main.py
can be used to benchmark a CoreML Model. It would calculate a precision x recall curve for every label and every folder of images.
See here for a detailed explanation.
First, you need some images with ground truth data to benchmark the model. The images can be in a nested folder structure to allow benchmarking categories of images. It's just important that you exactly mirror the folder structure with your ground truth data. Example structure:
- data
- images
- sharp_images
- black_white_images
- image1.jpg
- image2.jpg
- colored_images
- image3.jpg
- unsharp_images
- image4.jpg
- image5.jpg
- labels
- sharp_images
- black_white_images
- image1.txt
- image2.txt
- colored_images
- image3.txt
- unsharp_images
- image4.txt
- image5.txt
The script will then output detections for every image and one graph for each tag for each folder. So there will be one graph for all black_white_images, one for all colored_images, one for all sharp images ...
Furthermore, some values in the script should be changed according to your needs. In src/coreml_metrics/main.py
you'll find some global variables, which are specific to the concrete model you use:
classLabels
-> This is the list of labels your model recognized, all pretrained models of YOLOv5 used the coco dataset and therefore regonize 80 Labels
To run the script use the command:
$ poetry run coreml-metrics