🐙

PytorchのYOLOv5をCore MLフォーマットに変換する

2021/12/04に公開

Pytorchで書かれたYOLOv5をCore MLのフォーマットに変換する手順について紹介します。実はv6.0からYOLOv5のリリースには.mlmodelファイルもあります。これはYOLOv5のexport.pyからそのまま出てくるものとは違い、更にcoremltoolsを使った変換が必用です。なので、カスタムデータセットで学習して、自分の目的にあったモデルを作るにはyolov5-coreml-toolsでやっているような変換が必用になります。

coremltoolsの公式ドキュメントには、物体検出モデルの完全な変換や、MILに関する例がないので、YOLOv5とyolov5-coreml-toolsのソースコードをよく見ることで、さらに異なったモデルのCore MLへの変換にも役立つと思います。

環境

Python 3.7.10
Poetry version 1.1.12
Intel Mac (macOS 11.6)

エクスポート手順

$ ghq get https://github.com/dbsystel/yolov5-coreml-tools.git 
$ cd /path/to/yolov5-coreml-tools
$ mkdir models
$ cd models
$ curl -OL https://github.com/ultralytics/yolov5/releases/download/v4.0/yolov5s.pt
$ mv yolov5s.pt yolov5s_v4.pt
$ cd ../..
$ git clone https://github.com/ultralytics/yolov5.git
$ cd yolov5
$ git checkout refs/tags/v4.0

yolov5にsetup.pyを追加します。

setup.py

from setuptools import setup

setup(
    name="yolov5",
    url="https://github.com/ultralytics/yolov5",
    maintainer="ultralytics",
    maintainer_email="glenn.jocher@ultralytics.com",
    packages=["models", "utils"],
    install_requires=["opencv-python", "matplotlib", "torchvision", "PyYAML", "requests", "pandas", "seaborn"],
)

yolov5-coreml-toolsの方も1箇所修正します。ありがちなことですが、訓練モードのときと評価モードのときで戻り値の型が異なります。

src/coreml_export/main.py

@@ -312,7 +313,7 @@ def main():
     model = torch.load(opt.model_input_path, map_location=torch.device('cpu'))[
         'model'].float()
 
-    model.eval()
+    model.train()

準備が整ったのでpytorchのyolov5s_v4.ptを.mlmodelファイルに変換します。

$ cd /path/to/yolov5-coreml-tools
$ poetry install
$ poetry run coreml-export 
$ open models/yolov5-iOS.mlmodel #Xcodeで開く

最後にXCodeで動作確認してみます。

pixabay.com: 家族-ビーチ-人-海洋-6398107

写真は家族-ビーチ-人-海洋-6398107を使用。

環境

エクスポート手順

Discussion