diff --git a/tests/test_exports.py b/tests/test_exports.py index 76207230..cb734c3b 100644 --- a/tests/test_exports.py +++ b/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 diff --git a/ultralytics/engine/exporter.py b/ultralytics/engine/exporter.py index 227a76d5..f66f2be8 100644 --- a/ultralytics/engine/exporter.py +++ b/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)