Yolov5将pt权重转为coreml

简体中文 | English

介绍

本文提供了一个权重转换工具，可用于将 Yolov5 模型（例如，yolov5s.pt）导出到具有解码层和非最大抑制层(NMS)的 CoreML 模型(例如，yolov5s.mlmodel)。

轻松使用
在搭建好环境后只需要按照参数说明修改后，运行脚本就可以快速获取一些CoreML模型文件。

支持量化
支持FP16和Int8精度的转换，支持MacOs和Linux系统(Windows尚未测试)。

Xcode 预览
与官方的权重转换工具相比，我们添加了解码层和 NMS 层，以便导出的模型与 Apple 版本框架很好地配合，无需任何代码即可在 Xcode 上预览。

---

---

一、CoreML是什么？

Core ML是苹果公司于2017年6月宣布了上线的一种不需要互联网连接即可进行机器学习的框架。Core ML旨在优化在移动设备上训练机器学习的功效，且不需要互联网连接即可进行机器学习。^[1]
利用 Core ML 在 App 中整合机器学习模型。Core ML 为所有模型提供了一种统一的呈现方式。App 可以使用 Core ML API 和用户数据进行预测，以及训练或精调模型，一切都在用户设备上完成。^[2]

二、使用步骤

1.安装环境

运行此脚本需要yolov5 库(推荐yolov5-v6.2)，需要按照说明部署yolov5运行环境。

git clone https://github.com/ultralytics/yolov5  # clone
cd yolov5
git checkout v6.2                                # switch to v6.2 tag
pip install -r requirements.txt                  # install

运行脚本还需要以下coremltools库，安装命令如下

pip install coremltools==6.0

在linux的python3.8中部署yolov5 v6.2已经通过测试，其余情况自行测试。

2.开始使用

克隆脚本仓库

git clone https://github.com/ClintRen/yolov5_convert_weight_to_coreml.git

在convert.py可以看到如下参数和帮助，按照参数帮助来填写你的权重文件信息。
需要注意的是–yolov5-repo参数是必要的，它是yolov5目录的路径。

    parser.add_argument(
        "--yolov5-repo",
        required=True,
        # default="",
        help="path to yolov5 repo",
    )
    parser.add_argument(
        "--weight", type=str, default="yolov5s.pt", help="yolov5 weights path"
    )
    parser.add_argument(
        "--img-size",
        type=int,
        default=640,
        help="image imput size (pixels)",
    )
    parser.add_argument(
        "--conf-thres", type=float, default=0.25, help="confidence threshold"
    )
    parser.add_argument(
        "--iou-thres", type=float, default=0.45, help="NMS IoU threshold"
    )
    parser.add_argument(
        "--device", default="cpu", help="cuda device, i.e. 0 or 0,1,2,3 or cpu"
    )
    parser.add_argument(
        "--quantize",
        default=True,
        action="store_true",
        help="quantize model to FP16 and Int8",
    )

此外，以下额外信息会显示在CoreML文件的Metadata当中。

    # Other model options
    parser.add_argument("--description", default="yolov5s", help="Model description")
    parser.add_argument("--author", default="yolov5", help="Model author")
    parser.add_argument("--version", default="6.2", help="Model version")
    parser.add_argument("--license", default="GPL-3.0", help="Model license")

---

3.示例

在转换脚本目录下运行以下命令

python convert.py --yolov5-repo /path/to/yolov5 --weight yolov5s.pt --img-size 640 --quantize

命令说明

python convert.py             # 执行convert.py脚本
--yolov5-repo /path/to/yolov5 # yolov5路径
--weight yolov5s.pt           # 权重文件，本仓库自带yolov5s.pt供大家测试，它是yolov5-v6.2版本的
--img-size 640                # 输入图片尺寸
--quantize                    # 量化，额外输出FP16和Int8精度的CoreMl文件

运行完脚本后会在权重目录下生成3个CoreML文件，分别是yolov5s.mlmodel, yolov5s_FP16.mlmodel and yolov5s_Int8.mlmodel.

4.Xcode检测预览

将模型导入Xcode中，可以看到模型文件信息（如Metadata, Precision, Class Labels, Input, Output等）

下图是用yolov5s_Int8.mlmodel在Xcode-Preview中的的检测结果，无需编写额外的代码。

总结

以上就是快速将yolov5权重文件转为多精度且具有解码层和NMS层的CoreML文件和在Xcode预览的全过程。
如果对你有帮助麻烦你在本脚本仓库点一个star，谢谢。

参考

[1] CoreML报道

[2] CoreML官方