Philip Hoang

Data Scientist

2 minute read

If you have a trained detection model in PyTorch, you may want to export it to be deployed on an iOS device. To do this, you have to convert your saved model weights (.pt file) into a .mlmodel file.

The CoreMLTools Python library can be used to directly convert a PyTorch model into a CoreML model. The coremltools.convert function can trace the model and convert it directly. Unfortunately, this function skips out on many of the post-processing steps such as non-max suppression, the last sigmoid activation, and the conversion between cell-based coordinates into coordinates relative to the image.

Although there is an export function provided by Glenn Jocher and the YOLOv5 team, the trace function used in it does not export many of the post-processing steps such as adjusting the coordinates to be relative to the image rather than the grid cell. Fortunately, Leon de Andrade and Dennis Post have provided a repo to export the YOLOv5 model with all of these post-processing steps here.

We have used their repo to export our model with some minor modifications. Both the yolov5-coreml-tools repo and an older version of YOLOv5 (v4.0) need to be put into the same folder. The following instructions are based on the yolov5-coreml-tools repo. To export the model into CoreML, install poetry. Poetry is used to install the required libraries to export the model. Export it by navigating to the export/yolov5-coreml-tools folder and then using the following command

$ poetry install

CoreML tools only works for certain versions of PyTorch, so the following commands may be preferred

$ pyenv install 3.8.6
$ pyenv global 3.8.6
$ poetry install

In the src/coreml_export/main.py folder, you may want to change some of the variables such as

  • classLabels – A list of the names of the classes. In COCO, there are 80 classes.
  • anchor – This depends on the YOLOv5 model that you are using (x, m, l, xl). This can be found in the yolo.yml file
  • reverseModel – Some models reverse the order of the anchors and strides, so this is used to quickly reverse their order.

Then you may paste your .pt network model in the yolov5-coreml-tools folder. To finally run the export program, you may use the command

$ poetry run coreml-export --model-input-path <path to .pt file>

You can use the -h flag to get a list of the optional arguments for your export. The model will save your exported model in the output/models folder. If you would like to export a different model architecture, I suggest that you look at this post by Matthijs Hollemans or into the coremltools documentation.

You May Also Enjoy