I am done

rl/Lib/site-packages/onnx/backend/test/runner/__init__.py

@@ -0,0 +1,529 @@
+# Copyright (c) ONNX Project Contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+
+from __future__ import annotations
+
+import functools
+import glob
+import os
+import re
+import shutil
+import sys
+import tempfile
+import time
+import unittest
+from collections import defaultdict
+from typing import Any, Callable, Iterable, Pattern, Sequence
+from urllib.request import urlretrieve
+
+import numpy as np
+
+import onnx
+import onnx.reference
+from onnx import ONNX_ML, ModelProto, NodeProto, TypeProto, ValueInfoProto, numpy_helper
+from onnx.backend.base import Backend
+from onnx.backend.test.case.test_case import TestCase
+from onnx.backend.test.loader import load_model_tests
+from onnx.backend.test.runner.item import TestItem
+
+
+class BackendIsNotSupposedToImplementIt(unittest.SkipTest):
+    pass
+
+
+def retry_execute(times: int) -> Callable[[Callable[..., Any]], Callable[..., Any]]:
+    assert times >= 1
+
+    def wrapper(func: Callable[..., Any]) -> Callable[..., Any]:
+        @functools.wraps(func)
+        def wrapped(*args: Any, **kwargs: Any) -> Any:
+            for i in range(1, times + 1):
+                try:
+                    return func(*args, **kwargs)
+                except Exception:  # noqa: PERF203
+                    print(f"{i} times tried")
+                    if i == times:
+                        raise
+                    time.sleep(5 * i)
+
+        return wrapped
+
+    return wrapper
+
+
+class Runner:
+    def __init__(
+        self,
+        backend: type[Backend],
+        parent_module: str | None = None,
+        test_kwargs: dict | None = None,
+    ) -> None:
+        self.backend = backend
+        self._parent_module = parent_module
+        self._include_patterns: set[Pattern[str]] = set()
+        self._exclude_patterns: set[Pattern[str]] = set()
+        self._xfail_patterns: set[Pattern[str]] = set()
+        self._test_kwargs: dict = test_kwargs or {}
+
+        # This is the source of the truth of all test functions.
+        # Properties `test_cases`, `test_suite` and `tests` will be
+        # derived from it.
+        # {category: {name: func}}
+        self._test_items: dict[str, dict[str, TestItem]] = defaultdict(dict)
+
+        for rt in load_model_tests(kind="node"):
+            self._add_model_test(rt, "Node")
+
+        for rt in load_model_tests(kind="real"):
+            self._add_model_test(rt, "Real")
+
+        for rt in load_model_tests(kind="simple"):
+            self._add_model_test(rt, "Simple")
+
+        for ct in load_model_tests(kind="pytorch-converted"):
+            self._add_model_test(ct, "PyTorchConverted")
+
+        for ot in load_model_tests(kind="pytorch-operator"):
+            self._add_model_test(ot, "PyTorchOperator")
+
+    def _get_test_case(self, name: str) -> type[unittest.TestCase]:
+        test_case = type(str(name), (unittest.TestCase,), {})
+        if self._parent_module:
+            test_case.__module__ = self._parent_module
+        return test_case
+
+    def include(self, pattern: str) -> Runner:
+        self._include_patterns.add(re.compile(pattern))
+        return self
+
+    def exclude(self, pattern: str) -> Runner:
+        self._exclude_patterns.add(re.compile(pattern))
+        return self
+
+    def xfail(self, pattern: str) -> Runner:
+        self._xfail_patterns.add(re.compile(pattern))
+        return self
+
+    def enable_report(self) -> Runner:
+        import pytest
+
+        for category, items_map in self._test_items.items():
+            for item in items_map.values():
+                item.func = pytest.mark.onnx_coverage(item.proto, category)(item.func)
+        return self
+
+    @property
+    def _filtered_test_items(self) -> dict[str, dict[str, TestItem]]:
+        filtered: dict[str, dict[str, TestItem]] = {}
+        for category, items_map in self._test_items.items():
+            filtered[category] = {}
+            for name, item in items_map.items():
+                if self._include_patterns and (
+                    not any(include.search(name) for include in self._include_patterns)
+                ):
+                    item.func = unittest.skip("no matched include pattern")(item.func)
+                for exclude in self._exclude_patterns:
+                    if exclude.search(name):
+                        item.func = unittest.skip(
+                            f'matched exclude pattern "{exclude.pattern}"'
+                        )(item.func)
+                for xfail in self._xfail_patterns:
+                    if xfail.search(name):
+                        item.func = unittest.expectedFailure(item.func)
+                filtered[category][name] = item
+        return filtered
+
+    @property
+    def test_cases(self) -> dict[str, type[unittest.TestCase]]:
+        """List of test cases to be applied on the parent scope
+        Example usage:
+            globals().update(BackendTest(backend).test_cases)
+        """
+        test_cases = {}
+        for category, items_map in self._filtered_test_items.items():
+            test_case_name = f"OnnxBackend{category}Test"
+            test_case = self._get_test_case(test_case_name)
+            for name, item in sorted(items_map.items()):
+                setattr(test_case, name, item.func)
+            test_cases[test_case_name] = test_case
+        return test_cases
+
+    @property
+    def test_suite(self) -> unittest.TestSuite:
+        """TestSuite that can be run by TestRunner
+        Example usage:
+            unittest.TextTestRunner().run(BackendTest(backend).test_suite)
+        """
+        suite = unittest.TestSuite()
+        for case in sorted(
+            self.test_cases.values(), key=lambda cl: cl.__class__.__name__
+        ):
+            suite.addTests(unittest.defaultTestLoader.loadTestsFromTestCase(case))
+        return suite
+
+    # For backward compatibility (we used to expose `.tests`)
+    @property
+    def tests(self) -> type[unittest.TestCase]:
+        """One single unittest.TestCase that hosts all the test functions
+        Example usage:
+            onnx_backend_tests = BackendTest(backend).tests
+        """
+        tests = self._get_test_case("OnnxBackendTest")
+        for items_map in sorted(
+            self._filtered_test_items.values(), key=lambda cl: cl.__class__.__name__
+        ):
+            for name, item in sorted(items_map.items()):
+                setattr(tests, name, item.func)
+        return tests
+
+    @classmethod
+    def assert_similar_outputs(
+        cls,
+        ref_outputs: Sequence[Any],
+        outputs: Sequence[Any],
+        rtol: float,
+        atol: float,
+        model_dir: str | None = None,
+    ) -> None:
+        try:
+            np.testing.assert_equal(len(outputs), len(ref_outputs))
+        except TypeError as e:
+            raise TypeError(
+                f"Unable to compare expected type {type(ref_outputs)} "
+                f"and runtime type {type(outputs)} (known test={model_dir or '?'!r})"
+            ) from e
+        for i in range(len(outputs)):
+            if isinstance(outputs[i], (list, tuple)):
+                if not isinstance(ref_outputs[i], (list, tuple)):
+                    raise AssertionError(  # noqa: TRY004
+                        f"Unexpected type {type(outputs[i])} for outputs[{i}]. Expected "
+                        f"type is {type(ref_outputs[i])} (known test={model_dir or '?'!r})."
+                    )
+                for j in range(len(outputs[i])):
+                    cls.assert_similar_outputs(
+                        ref_outputs[i][j],
+                        outputs[i][j],
+                        rtol,
+                        atol,
+                        model_dir=model_dir,
+                    )
+            else:
+                if not np.issubdtype(ref_outputs[i].dtype, np.number):
+                    if ref_outputs[i].tolist() != outputs[i].tolist():
+                        raise AssertionError(f"{ref_outputs[i]} != {outputs[i]}")
+                    continue
+                np.testing.assert_equal(outputs[i].dtype, ref_outputs[i].dtype)
+                if ref_outputs[i].dtype == object:  # type: ignore[attr-defined]
+                    np.testing.assert_array_equal(outputs[i], ref_outputs[i])
+                else:
+                    np.testing.assert_allclose(
+                        outputs[i], ref_outputs[i], rtol=rtol, atol=atol
+                    )
+
+    @classmethod
+    @retry_execute(3)
+    def download_model(
+        cls,
+        model_test: TestCase,
+        model_dir: str,
+        models_dir: str,
+    ) -> None:
+        # On Windows, NamedTemporaryFile can not be opened for a
+        # second time
+        del model_dir
+        download_file = tempfile.NamedTemporaryFile(delete=False)
+        try:
+            download_file.close()
+            assert model_test.url
+            print(
+                f"Start downloading model {model_test.model_name} from {model_test.url}"
+            )
+            urlretrieve(model_test.url, download_file.name)
+            print("Done")
+            onnx.utils._extract_model_safe(download_file.name, models_dir)
+        except Exception as e:
+            print(f"Failed to prepare data for model {model_test.model_name}: {e}")
+            raise
+        finally:
+            os.remove(download_file.name)
+
+    @classmethod
+    def prepare_model_data(cls, model_test: TestCase) -> str:
+        onnx_home = os.path.expanduser(
+            os.getenv("ONNX_HOME", os.path.join("~", ".onnx"))
+        )
+        models_dir = os.getenv("ONNX_MODELS", os.path.join(onnx_home, "models"))
+        model_dir: str = os.path.join(models_dir, model_test.model_name)
+        if not os.path.exists(os.path.join(model_dir, "model.onnx")):
+            if os.path.exists(model_dir):
+                bi = 0
+                while True:
+                    dest = f"{model_dir}.old.{bi}"
+                    if os.path.exists(dest):
+                        bi += 1
+                        continue
+                    shutil.move(model_dir, dest)
+                    break
+            os.makedirs(model_dir)
+
+            cls.download_model(
+                model_test=model_test, model_dir=model_dir, models_dir=models_dir
+            )
+        return model_dir
+
+    def _add_test(
+        self,
+        category: str,
+        test_name: str,
+        test_func: Callable[..., Any],
+        report_item: list[ModelProto | NodeProto | None],
+        devices: Iterable[str] = ("CPU", "CUDA"),
+        **kwargs: Any,
+    ) -> None:
+        # We don't prepend the 'test_' prefix to improve greppability
+        if not test_name.startswith("test_"):
+            raise ValueError(f"Test name must start with test_: {test_name}")
+
+        def add_device_test(device: str) -> None:
+            device_test_name = f"{test_name}_{device.lower()}"
+            if device_test_name in self._test_items[category]:
+                raise ValueError(
+                    f'Duplicated test name "{device_test_name}" in category "{category}"'
+                )
+
+            @unittest.skipIf(  # type: ignore
+                not self.backend.supports_device(device),
+                f"Backend doesn't support device {device}",
+            )
+            @functools.wraps(test_func)
+            def device_test_func(*args: Any, **device_test_kwarg: Any) -> Any:
+                try:
+                    merged_kwargs = {**kwargs, **device_test_kwarg}
+                    return test_func(*args, device, **merged_kwargs)
+                except BackendIsNotSupposedToImplementIt as e:
+                    # hacky verbose reporting
+                    if "-v" in sys.argv or "--verbose" in sys.argv:
+                        print(f"Test {device_test_name} is effectively skipped: {e}")
+
+            self._test_items[category][device_test_name] = TestItem(
+                device_test_func, report_item
+            )
+
+        for device in devices:
+            add_device_test(device)
+
+    @staticmethod
+    def generate_dummy_data(
+        x: ValueInfoProto, seed: int = 0, name: str = "", random: bool = False
+    ) -> np.ndarray:
+        """Generates a random tensor based on the input definition."""
+        if not x.type.tensor_type:
+            raise NotImplementedError(
+                f"Input expected to have tensor type. "
+                f"Unable to generate random data for model {name!r} and input {x}."
+            )
+        if x.type.tensor_type.elem_type != 1:
+            raise NotImplementedError(
+                f"Currently limited to float tensors. "
+                f"Unable to generate random data for model {name!r} and input {x}."
+            )
+        shape = tuple(
+            d.dim_value if d.HasField("dim_value") else 1
+            for d in x.type.tensor_type.shape.dim
+        )
+        if random:
+            gen = np.random.default_rng(seed=seed)
+            return gen.random(shape, np.float32)
+        n = np.prod(shape)
+        return (np.arange(n).reshape(shape) / n).astype(np.float32)
+
+    def _add_model_test(self, model_test: TestCase, kind: str) -> None:
+        # model is loaded at runtime, note sometimes it could even
+        # never loaded if the test skipped
+        model_marker: list[ModelProto | NodeProto | None] = [None]
+
+        def run(test_self: Any, device: str, **kwargs) -> None:  # noqa: ARG001
+            if model_test.url is not None and model_test.url.startswith(
+                "onnx/backend/test/data/light/"
+            ):
+                # testing local files
+                model_pb_path = os.path.normpath(
+                    os.path.join(
+                        os.path.dirname(__file__),
+                        "..",
+                        "..",
+                        "..",
+                        "..",
+                        model_test.url,
+                    )
+                )
+                if not os.path.exists(model_pb_path):
+                    raise FileNotFoundError(f"Unable to find model {model_pb_path!r}.")
+                onnx_home = os.path.expanduser(
+                    os.getenv("ONNX_HOME", os.path.join("~", ".onnx"))
+                )
+                models_dir = os.getenv(
+                    "ONNX_MODELS", os.path.join(onnx_home, "models", "light")
+                )
+                model_dir: str = os.path.join(models_dir, model_test.model_name)
+                if not os.path.exists(model_dir):
+                    os.makedirs(model_dir)
+                use_dummy = True
+            else:
+                if model_test.model_dir is None:
+                    model_dir = self.prepare_model_data(model_test)
+                else:
+                    model_dir = model_test.model_dir
+                model_pb_path = os.path.join(model_dir, "model.onnx")
+                use_dummy = False
+
+            if not ONNX_ML and "ai_onnx_ml" in model_dir:
+                return
+
+            model = onnx.load(model_pb_path)
+            model_marker[0] = model
+            if (
+                hasattr(self.backend, "is_compatible")
+                and callable(self.backend.is_compatible)
+                and not self.backend.is_compatible(model)
+            ):
+                raise unittest.SkipTest("Not compatible with backend")
+
+            prepared_model = self.backend.prepare(model, device, **kwargs)
+            assert prepared_model is not None
+
+            if use_dummy:
+                # When the backend test goes through a test involving a
+                # model stored in onnx/backend/test/data/light,
+                # this function generates expected output coming from
+                # from ReferenceEvaluator run with random inputs.
+                # A couple of models include many Conv operators and the
+                # python implementation is slow (such as test_bvlc_alexnet).
+                with open(model_pb_path, "rb") as f:
+                    onx = onnx.load(f)
+
+                test_data_set = os.path.join(model_dir, "test_data_set_0")
+                if not os.path.exists(test_data_set):
+                    os.mkdir(test_data_set)
+                feeds = {}
+                inits = {i.name for i in onx.graph.initializer}
+                n_input = 0
+                inputs = []
+                for i in range(len(onx.graph.input)):
+                    if onx.graph.input[i].name in inits:
+                        continue
+                    name = os.path.join(test_data_set, f"input_{n_input}.pb")
+                    inputs.append(name)
+                    n_input += 1
+                    x = onx.graph.input[i]
+                    value = self.generate_dummy_data(
+                        x, seed=0, name=model_test.model_name, random=False
+                    )
+                    feeds[x.name] = value
+                    with open(name, "wb") as f:
+                        f.write(onnx.numpy_helper.from_array(value).SerializeToString())
+
+                # loads expected output if any available
+                prefix = os.path.splitext(model_pb_path)[0]
+                expected_outputs = []
+                for i in range(len(onx.graph.output)):
+                    name = f"{prefix}_output_{i}.pb"
+                    if os.path.exists(name):
+                        expected_outputs.append(name)
+                        continue
+                    expected_outputs = None
+                    break
+
+                if expected_outputs is None:
+                    ref = onnx.reference.ReferenceEvaluator(onx)
+                    outputs = ref.run(None, feeds)
+                    for i, o in enumerate(outputs):
+                        name = os.path.join(test_data_set, f"output_{i}.pb")
+                        with open(name, "wb") as f:
+                            f.write(onnx.numpy_helper.from_array(o).SerializeToString())
+                else:
+                    for i, o in enumerate(expected_outputs):
+                        name = os.path.join(test_data_set, f"output_{i}.pb")
+                        shutil.copy(o, name)
+            else:
+                # TODO after converting all npz files to protobuf, we can delete this.
+                for test_data_npz in glob.glob(
+                    os.path.join(model_dir, "test_data_*.npz")
+                ):
+                    test_data = np.load(test_data_npz, encoding="bytes")
+                    inputs = list(test_data["inputs"])
+                    outputs = list(prepared_model.run(inputs))
+                    ref_outputs = tuple(
+                        np.array(x) if not isinstance(x, (list, dict)) else x
+                        for f in test_data["outputs"]
+                    )
+                    self.assert_similar_outputs(
+                        ref_outputs,
+                        outputs,
+                        rtol=kwargs.get("rtol", model_test.rtol),
+                        atol=kwargs.get("atol", model_test.atol),
+                        model_dir=model_dir,
+                    )
+
+            for test_data_dir in glob.glob(os.path.join(model_dir, "test_data_set*")):
+                inputs = []
+                inputs_num = len(glob.glob(os.path.join(test_data_dir, "input_*.pb")))
+                for i in range(inputs_num):
+                    input_file = os.path.join(test_data_dir, f"input_{i}.pb")
+                    self._load_proto(input_file, inputs, model.graph.input[i].type)
+                ref_outputs = []
+                ref_outputs_num = len(
+                    glob.glob(os.path.join(test_data_dir, "output_*.pb"))
+                )
+                for i in range(ref_outputs_num):
+                    output_file = os.path.join(test_data_dir, f"output_{i}.pb")
+                    self._load_proto(
+                        output_file, ref_outputs, model.graph.output[i].type
+                    )
+                outputs = list(prepared_model.run(inputs))
+                self.assert_similar_outputs(
+                    ref_outputs,
+                    outputs,
+                    rtol=kwargs.get("rtol", model_test.rtol),
+                    atol=kwargs.get("atol", model_test.atol),
+                    model_dir=model_dir,
+                )
+
+        if model_test.name in self._test_kwargs:
+            self._add_test(
+                kind + "Model",
+                model_test.name,
+                run,
+                model_marker,
+                **self._test_kwargs[model_test.name],
+            )
+        else:
+            self._add_test(kind + "Model", model_test.name, run, model_marker)
+
+    def _load_proto(
+        self,
+        proto_filename: str,
+        target_list: list[np.ndarray | list[Any]],
+        model_type_proto: TypeProto,
+    ) -> None:
+        with open(proto_filename, "rb") as f:
+            protobuf_content = f.read()
+            if model_type_proto.HasField("sequence_type"):
+                sequence = onnx.SequenceProto()
+                sequence.ParseFromString(protobuf_content)
+                target_list.append(numpy_helper.to_list(sequence))
+            elif model_type_proto.HasField("tensor_type"):
+                tensor = onnx.TensorProto()
+                tensor.ParseFromString(protobuf_content)
+                t = numpy_helper.to_array(tensor)
+                assert isinstance(t, np.ndarray)
+                target_list.append(t)
+            elif model_type_proto.HasField("optional_type"):
+                optional = onnx.OptionalProto()
+                optional.ParseFromString(protobuf_content)
+                target_list.append(numpy_helper.to_optional(optional))  # type: ignore[arg-type]
+            else:
+                print(
+                    "Loading proto of that specific type (Map/Sparse Tensor) is currently not supported"
+                )

rl/Lib/site-packages/onnx/backend/test/runner/item.py

@@ -0,0 +1,18 @@
+# Copyright (c) ONNX Project Contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+from __future__ import annotations
+
+import dataclasses
+from typing import Any, Callable
+
+from onnx import ModelProto, NodeProto
+
+# A container that hosts the test function and the associated
+# test item (ModelProto)
+
+
+@dataclasses.dataclass
+class TestItem:
+    func: Callable[..., Any]
+    proto: list[ModelProto | NodeProto | None]