New ONNX, TorchScript, CoreML export test matrices (#11652)

Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>

tests/test_exports.py

@@ -20,22 +20,15 @@ from ultralytics.utils import (
 from ultralytics.utils.torch_utils import TORCH_1_9, TORCH_1_13
 from . import MODEL, SOURCE
 
-# Constants
-EXPORT_PARAMETERS_LIST = [  # generate all combinations but exclude those where both int8 and half are True
-    (task, dynamic, int8, half, batch)
-    for task, dynamic, int8, half, batch in product(TASKS, [True, False], [True, False], [True, False], [1, 2])
-    if not (int8 and half)  # exclude cases where both int8 and half are True
-]
-
 
 def test_export_torchscript():
-    """Test exporting the YOLO model to TorchScript format."""
+    """Test YOLO exports to TorchScript format."""
     f = YOLO(MODEL).export(format="torchscript", optimize=False, imgsz=32)
     YOLO(f)(SOURCE, imgsz=32)  # exported model inference
 
 
 def test_export_onnx():
-    """Test exporting the YOLO model to ONNX format."""
+    """Test YOLO exports to ONNX format."""
     f = YOLO(MODEL).export(format="onnx", dynamic=True, imgsz=32)
     YOLO(f)(SOURCE, imgsz=32)  # exported model inference
 
@@ -43,7 +36,7 @@ def test_export_onnx():
 @pytest.mark.skipif(checks.IS_PYTHON_3_12, reason="OpenVINO not supported in Python 3.12")
 @pytest.mark.skipif(not TORCH_1_13, reason="OpenVINO requires torch>=1.13")
 def test_export_openvino():
-    """Test exporting the YOLO model to OpenVINO format."""
+    """Test YOLO exports to OpenVINO format."""
     f = YOLO(MODEL).export(format="openvino", imgsz=32)
     YOLO(f)(SOURCE, imgsz=32)  # exported model inference
 
@@ -51,9 +44,16 @@ def test_export_openvino():
 @pytest.mark.slow
 @pytest.mark.skipif(checks.IS_PYTHON_3_12, reason="OpenVINO not supported in Python 3.12")
 @pytest.mark.skipif(not TORCH_1_13, reason="OpenVINO requires torch>=1.13")
-@pytest.mark.parametrize("task, dynamic, int8, half, batch", EXPORT_PARAMETERS_LIST)
+@pytest.mark.parametrize(
+    "task, dynamic, int8, half, batch",
+    [  # generate all combinations but exclude those where both int8 and half are True
+        (task, dynamic, int8, half, batch)
+        for task, dynamic, int8, half, batch in product(TASKS, [True, False], [True, False], [True, False], [1, 2])
+        if not (int8 and half)  # exclude cases where both int8 and half are True
+    ],
+)
 def test_export_openvino_matrix(task, dynamic, int8, half, batch):
-    """Test exporting the YOLO model to OpenVINO format."""
+    """Test YOLO exports to OpenVINO format."""
     file = YOLO(TASK2MODEL[task]).export(
         format="openvino",
         imgsz=32,
@@ -73,12 +73,68 @@ def test_export_openvino_matrix(task, dynamic, int8, half, batch):
         shutil.rmtree(file)
 
 
+@pytest.mark.slow
+@pytest.mark.parametrize("task, dynamic, int8, half, batch", product(TASKS, [True, False], [False], [False], [1, 2]))
+def test_export_onnx_matrix(task, dynamic, int8, half, batch):
+    """Test YOLO exports to ONNX format."""
+    file = YOLO(TASK2MODEL[task]).export(
+        format="onnx",
+        imgsz=32,
+        dynamic=dynamic,
+        int8=int8,
+        half=half,
+        batch=batch,
+    )
+    YOLO(file)([SOURCE] * batch, imgsz=64 if dynamic else 32)  # exported model inference
+    Path(file).unlink()  # cleanup
+
+
+@pytest.mark.slow
+@pytest.mark.parametrize("task, dynamic, int8, half, batch", product(TASKS, [False], [False], [False], [1, 2]))
+def test_export_torchscript_matrix(task, dynamic, int8, half, batch):
+    """Test YOLO exports to TorchScript format."""
+    file = YOLO(TASK2MODEL[task]).export(
+        format="torchscript",
+        imgsz=32,
+        dynamic=dynamic,
+        int8=int8,
+        half=half,
+        batch=batch,
+    )
+    YOLO(file)([SOURCE] * 3, imgsz=64 if dynamic else 32)  # exported model inference at batch=3
+    Path(file).unlink()  # cleanup
+
+
+@pytest.mark.slow
+@pytest.mark.skipif(not MACOS, reason="CoreML inference only supported on macOS")
+@pytest.mark.parametrize(
+    "task, dynamic, int8, half, batch",
+    [  # generate all combinations but exclude those where both int8 and half are True
+        (task, dynamic, int8, half, batch)
+        for task, dynamic, int8, half, batch in product(TASKS, [False], [True, False], [True, False], [1])
+        if not (int8 and half)  # exclude cases where both int8 and half are True
+    ],
+)
+def test_export_coreml_matrix(task, dynamic, int8, half, batch):
+    """Test YOLO exports to TorchScript format."""
+    file = YOLO(TASK2MODEL[task]).export(
+        format="coreml",
+        imgsz=32,
+        dynamic=dynamic,
+        int8=int8,
+        half=half,
+        batch=batch,
+    )
+    YOLO(file)([SOURCE] * batch, imgsz=32)  # exported model inference at batch=3
+    shutil.rmtree(file)  # cleanup
+
+
 @pytest.mark.skipif(not TORCH_1_9, reason="CoreML>=7.2 not supported with PyTorch<=1.8")
 @pytest.mark.skipif(WINDOWS, reason="CoreML not supported on Windows")  # RuntimeError: BlobWriter not loaded
 @pytest.mark.skipif(IS_RASPBERRYPI, reason="CoreML not supported on Raspberry Pi")
 @pytest.mark.skipif(checks.IS_PYTHON_3_12, reason="CoreML not supported in Python 3.12")
 def test_export_coreml():
-    """Test exporting the YOLO model to CoreML format."""
+    """Test YOLO exports to CoreML format."""
     if MACOS:
         f = YOLO(MODEL).export(format="coreml", imgsz=32)
         YOLO(f)(SOURCE, imgsz=32)  # model prediction only supported on macOS for nms=False models
@@ -89,7 +145,7 @@ def test_export_coreml():
 @pytest.mark.skipif(not LINUX, reason="Test disabled as TF suffers from install conflicts on Windows and macOS")
 def test_export_tflite():
     """
-    Test exporting the YOLO model to TFLite format.
+    Test YOLO exports to TFLite format.
 
     Note TF suffers from install conflicts on Windows and macOS.
     """
@@ -102,7 +158,7 @@ def test_export_tflite():
 @pytest.mark.skipif(not LINUX, reason="TF suffers from install conflicts on Windows and macOS")
 def test_export_pb():
     """
-    Test exporting the YOLO model to *.pb format.
+    Test YOLO exports to *.pb format.
 
     Note TF suffers from install conflicts on Windows and macOS.
     """
@@ -114,7 +170,7 @@ def test_export_pb():
 @pytest.mark.skipif(True, reason="Test disabled as Paddle protobuf and ONNX protobuf requirementsk conflict.")
 def test_export_paddle():
     """
-    Test exporting the YOLO model to Paddle format.
+    Test YOLO exports to Paddle format.
 
     Note Paddle protobuf requirements conflicting with onnx protobuf requirements.
     """
@@ -123,6 +179,6 @@ def test_export_paddle():
 
 @pytest.mark.slow
 def test_export_ncnn():
-    """Test exporting the YOLO model to NCNN format."""
+    """Test YOLO exports to NCNN format."""
     f = YOLO(MODEL).export(format="ncnn", imgsz=32)
     YOLO(f)(SOURCE, imgsz=32)  # exported model inference

ultralytics/engine/exporter.py

@@ -599,6 +599,7 @@ class Exporter:
 
         LOGGER.info(f"\n{prefix} starting export with coremltools {ct.__version__}...")
         assert not WINDOWS, "CoreML export is not supported on Windows, please run on macOS or Linux."
+        assert self.args.batch == 1, "CoreML batch sizes > 1 are not supported. Please retry at 'batch=1'."
         f = self.file.with_suffix(".mlmodel" if mlmodel else ".mlpackage")
         if f.is_dir():
             shutil.rmtree(f)