MITライセンスのYolo v9がCoreMLで動くそうなので試してみた

まずはiPhoneシュミレータでのYolo-v9-sの実行結果です。

Issue46でONNXへの変換サポートについて議論があり、そこでCoreMLへの変換についても話が上がっていました。

既にPRも出来ていましたので試してみたところエラーが出て動きません。
少し触ってみて何とか変換ができたのですが、NMSが組み込まれていないため、このままでは利用ができない状況でした。

まぁここまで動いていれば後はなんとかなるだろうと試行錯誤してみました。

というわけで、修正したコードが下記になります。
このパッチを組み込んで実行すると、NMSで処理した結果まで返すモデルが出来上がります。

diff --git a/yolo/tools/solver.py b/yolo/tools/solver.py
index 03bbcd2..545a86f 100644
--- a/yolo/tools/solver.py
+++ b/yolo/tools/solver.py
@@ -169,13 +169,13 @@ class ExportModel(BaseModel):
 
         export_mode = False
         format = cfg.task.format
+        self.format = format
         # TODO check if we can use export mode for all formats
         if self.format == "coreml":
             export_mode = True
 
         super().__init__(cfg, export_mode=export_mode)
         self.cfg = cfg
-        self.format = format
         self.model_exporter = ModelExporter(self.cfg, self.model, format=self.format)
 
     def export(self):
diff --git a/yolo/utils/export_utils.py b/yolo/utils/export_utils.py
index 906bc56..380eedb 100644
--- a/yolo/utils/export_utils.py
+++ b/yolo/utils/export_utils.py
@@ -82,20 +82,45 @@ class ModelExporter:
 
         self.model.eval()
         example_inputs = (torch.rand(1, 3, *self.cfg.image_size),)
-        exported_program = torch.export.export(self.model, example_inputs)
 
+        nms_cfg = NMSConfig(0.1, 0.5, 300)
+        model = CoremlPostProcess(self.model, nms_cfg)
+        exported_program = torch.jit.trace(model, example_inputs, strict=True)
         import logging
 
         import coremltools as ct
 
         # Convert to Core ML program using the Unified Conversion API.
         logging.getLogger("coremltools").disabled = True
-
         self.output_names: List[str] = ["preds_cls", "preds_anc", "preds_box"]
-
         model_from_export = ct.convert(
-            exported_program, outputs=[ct.TensorType(name=name) for name in self.output_names], convert_to="mlprogram"
+            exported_program,
+            inputs=[ct.ImageType("image", shape=example_inputs[0].shape,
+                                 scale=1/255., bias=[0,0,0])],
+            outputs=[ct.TensorType("predict")],
+            convert_to="mlprogram",
+            minimum_deployment_target=ct.target.iOS15,
+            compute_precision=ct.precision.FLOAT16,
         )
 
         model_from_export.save(self.model_path)
         logger.info(f":white_check_mark: Model exported to coreml format {self.model_path}")
+
+from torch import nn
+from yolo.utils.bounding_box_utils import bbox_nms
+from yolo.config.config import NMSConfig
+
+class CoremlPostProcess(nn.Module):
+
+    def __init__(self, model, nms_cfg):
+        super().__init__()
+        self.model = model
+        self.nms_cfg = nms_cfg
+
+    def forward(self, x):
+        prediction = self.model(x)
+
+        pred_class, _, pred_bbox = prediction[:3]
+        pred_conf = prediction[3] if len(prediction) == 4 else None
+        pred_bbox = bbox_nms(pred_class, pred_bbox, self.nms_cfg, pred_conf)
+        return pred_bbox
\ No newline at end of file

Pythonでの動作確認コードとSwiftのコードについてはこちらを参照ください。
変換したモデルも一緒に置いています。

最適化

パフォーマンスが良くないということで、パフォーマンス計測できる書き方に変更。
GPUで動作するようになりました。
ANEで動かすにはもう少し頑張る必要がありそうです。
遅かった理由としてはnmsをcoreml用のnmsに変更していなかったのが原因と思われます。

修正したコードとモデルはこちら

変換に利用した元のコードはこちら