From a892215f1a1c1dd113c426dcb0e824776666e51b Mon Sep 17 00:00:00 2001 From: Adrian Boguszewski Date: Tue, 13 Feb 2024 22:17:19 +0100 Subject: [PATCH] Test OpenVINO export (#1681) Signed-off-by: Adrian Boguszewski Co-authored-by: Samet Akcay --- tests/integration/model/test_models.py | 17 +++++++++++++++-- 1 file changed, 15 insertions(+), 2 deletions(-) diff --git a/tests/integration/model/test_models.py b/tests/integration/model/test_models.py index 3dc2d7880d..27788f506e 100644 --- a/tests/integration/model/test_models.py +++ b/tests/integration/model/test_models.py @@ -24,6 +24,11 @@ def models() -> list[str]: return get_available_models() +def export_types() -> list[ExportType]: + """Return all available export frameworks.""" + return list(ExportType) + + class TestAPI: """Do sanity check on all models.""" @@ -116,11 +121,19 @@ def test_predict(self, model_name: str, dataset_path: Path, project_path: Path) ) @pytest.mark.parametrize("model_name", models()) - def test_export(self, model_name: str, dataset_path: Path, project_path: Path) -> None: + @pytest.mark.parametrize("export_type", export_types()) + def test_export( + self, + model_name: str, + export_type: ExportType, + dataset_path: Path, + project_path: Path, + ) -> None: """Export model from checkpoint. Args: model_name (str): Name of the model. + export_type (ExportType): Framework to export to. dataset_path (Path): Root to dataset from fixture. project_path (Path): Path to temporary project folder from fixture. """ @@ -145,7 +158,7 @@ def test_export(self, model_name: str, dataset_path: Path, project_path: Path) - model=model, datamodule=dataset, ckpt_path=f"{project_path}/{model_name}/dummy/weights/last.ckpt", - export_type=ExportType.ONNX, + export_type=export_type, input_size=input_size, )