HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

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

Browse source 
Signed-off-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Glenn Jocher 2024-05-05 16:57:52 +02:00 • committed by GitHub
parent 2583f842b8
commit 972f615110
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 74 additions and 17 deletions
Download patch file Download diff file Expand all files Collapse all files

90
tests/test_exports.py
View file

@ -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