|
| 1 | +import sys |
| 2 | + |
| 3 | +import pytest |
| 4 | +import torch |
| 5 | + |
| 6 | +pytest.importorskip("executorch.runtime") |
| 7 | + |
| 8 | +from executorch.backends.xnnpack.partition.xnnpack_partitioner import XnnpackPartitioner # noqa: E402 |
| 9 | +from executorch.exir import to_edge_transform_and_lower # noqa: E402 |
| 10 | + |
| 11 | +from sab.runtimes.executorch import ExecuTorchRuntime # noqa: E402 |
| 12 | + |
| 13 | +pytestmark = pytest.mark.executorch |
| 14 | + |
| 15 | + |
| 16 | +class ConvModel(torch.nn.Module): |
| 17 | + """One NCHW image input.""" |
| 18 | + |
| 19 | + def __init__(self): |
| 20 | + super().__init__() |
| 21 | + self.conv = torch.nn.Conv2d(3, 4, 3, padding=1) |
| 22 | + |
| 23 | + def forward(self, images): |
| 24 | + return self.conv(images) |
| 25 | + |
| 26 | + |
| 27 | +class ImageAndSizesModel(ConvModel): |
| 28 | + """Image first, then an int64 input, like the `orig_target_sizes` input of some detectors.""" |
| 29 | + |
| 30 | + def forward(self, images, sizes): |
| 31 | + return self.conv(images), sizes + 1 |
| 32 | + |
| 33 | + |
| 34 | +class SizesAndImageModel(ConvModel): |
| 35 | + def forward(self, sizes, images): |
| 36 | + return sizes + 1, self.conv(images) |
| 37 | + |
| 38 | + |
| 39 | +CASES = { |
| 40 | + "conv": (ConvModel, (torch.rand(1, 3, 8, 8),)), |
| 41 | + "image_and_sizes": (ImageAndSizesModel, (torch.rand(1, 3, 8, 8), torch.tensor([[10, 20]]))), |
| 42 | + "sizes_and_image": (SizesAndImageModel, (torch.tensor([[10, 20]]), torch.rand(1, 3, 8, 8))), |
| 43 | +} |
| 44 | + |
| 45 | + |
| 46 | +@pytest.fixture(scope="session") |
| 47 | +def model_dir(tmp_path_factory): |
| 48 | + return tmp_path_factory.mktemp("pte") |
| 49 | + |
| 50 | + |
| 51 | +@pytest.fixture(scope="session") |
| 52 | +def build_pte(model_dir): |
| 53 | + """Export each tiny model once per session. Returns (module, example inputs, path).""" |
| 54 | + built = {} |
| 55 | + |
| 56 | + def build(case: str): |
| 57 | + if case not in built: |
| 58 | + torch.manual_seed(0) |
| 59 | + model_class, example_inputs = CASES[case] |
| 60 | + module = model_class().eval() |
| 61 | + exported = torch.export.export(module, example_inputs) |
| 62 | + program = to_edge_transform_and_lower(exported, partitioner=[XnnpackPartitioner()]).to_executorch() |
| 63 | + path = model_dir / f"{case}.pte" |
| 64 | + path.write_bytes(program.buffer) |
| 65 | + built[case] = (module, example_inputs, str(path)) |
| 66 | + return built[case] |
| 67 | + |
| 68 | + return build |
| 69 | + |
| 70 | + |
| 71 | +@pytest.mark.parametrize("platform", ["linux", "darwin"]) |
| 72 | +def test_cpu_is_available_on_linux_and_macos(monkeypatch, platform): |
| 73 | + monkeypatch.setattr(sys, "platform", platform) |
| 74 | + assert ExecuTorchRuntime.is_available("cpu") |
| 75 | + |
| 76 | + |
| 77 | +def test_run_matches_the_torch_model_and_outputs_own_their_memory(build_pte): |
| 78 | + module, _, path = build_pte("conv") |
| 79 | + runtime = ExecuTorchRuntime(path, "cpu", "fp32") |
| 80 | + images = torch.rand(1, 3, 8, 8) |
| 81 | + |
| 82 | + first = runtime.run({"input0": images}) |
| 83 | + runtime.run({"input0": torch.rand(1, 3, 8, 8)}) |
| 84 | + |
| 85 | + assert runtime.input_spec.name == "input0" |
| 86 | + assert set(first) == {"output0"} |
| 87 | + torch.testing.assert_close(first["output0"], module(images), atol=1e-4, rtol=1e-4) |
| 88 | + |
| 89 | + |
| 90 | +def test_run_passes_inputs_in_positional_order(build_pte): |
| 91 | + module, _, path = build_pte("image_and_sizes") |
| 92 | + runtime = ExecuTorchRuntime(path, "cpu", "fp32") |
| 93 | + images = torch.rand(1, 3, 8, 8) |
| 94 | + sizes = torch.tensor([[10, 20]]) |
| 95 | + |
| 96 | + # The dict order is the reverse of the positional order. |
| 97 | + outputs = runtime.run({"input1": sizes, "input0": images}) |
| 98 | + |
| 99 | + assert runtime.input_spec.name == "input0" |
| 100 | + torch.testing.assert_close(outputs["output0"], module(images, sizes)[0], atol=1e-4, rtol=1e-4) |
| 101 | + assert outputs["output1"].tolist() == [[11.0, 21.0]] |
| 102 | + |
| 103 | + |
| 104 | +def test_run_finds_the_image_at_a_later_position(build_pte): |
| 105 | + module, _, path = build_pte("sizes_and_image") |
| 106 | + runtime = ExecuTorchRuntime(path, "cpu", "fp32") |
| 107 | + images = torch.rand(1, 3, 8, 8) |
| 108 | + sizes = torch.tensor([[10, 20]]) |
| 109 | + |
| 110 | + outputs = runtime.run({"input0": sizes, "input1": images}) |
| 111 | + |
| 112 | + assert runtime.input_spec.name == "input1" |
| 113 | + assert runtime.input_spec.shape == (1, 3, 8, 8) |
| 114 | + assert outputs["output0"].tolist() == [[11.0, 21.0]] |
| 115 | + torch.testing.assert_close(outputs["output1"], module(sizes, images)[1], atol=1e-4, rtol=1e-4) |
| 116 | + |
| 117 | + |
| 118 | +def test_run_records_one_timing_per_call_after_warmup(build_pte): |
| 119 | + *_, path = build_pte("conv") |
| 120 | + runtime = ExecuTorchRuntime(path, "cpu", "fp32") |
| 121 | + assert runtime.profiler.timings == [] |
| 122 | + |
| 123 | + for expected in (1, 2, 3): |
| 124 | + runtime.run({"input0": torch.rand(1, 3, 8, 8)}) |
| 125 | + assert len(runtime.profiler.timings) == expected |
0 commit comments