Aidge interop torch API ======================== .. currentmodule:: aidge_interop_torch .. contents:: Table of Contents :depth: 3 :local: Introduction ------------ ``aidge_interop_torch`` provides interoperability between Aidge and PyTorch by wrapping a ``torch.nn.Module`` into an Aidge graph (via ONNX export/import) so that the two frameworks can run side by side on the same weights and inputs. The test suite documented here does not test business features — it validates the **numerical and performance equivalence** between a PyTorch model and its Aidge counterpart, for both the forward pass and the backward pass (gradients and loss). Every test in the suite follows the same pattern: 1. Build a PyTorch model. 2. Convert it into an Aidge model with :func:`aidge_interop_torch.wrap`. 3. Run identical inputs through both models. 4. Compare outputs, loss, gradients and (optionally) wall-clock time. The :func:`aidge_interop_torch.wrap` function exports the PyTorch model to ONNX, imports it into Aidge, and copies the model parameters. This ensures that both frameworks execute the same computation graph with identical weights, making numerical comparisons meaningful. Basic example of a comparison, for a single step (see :class:`ModelTester` for the full multi-epoch driver): .. code-block:: python import torch import aidge_interop_torch from utils import ModelTester torch_model = torch.nn.Linear(16, 4) aidge_model = aidge_interop_torch.wrap(torch_model, [(1, 16)]) tester = ModelTester( model1=torch_model, model2=aidge_model, name="Linear_example", test_backward=True, epochs=10, ) tester.test_multiple_step(input_sizes=(1, 16), label_sizes=(1, 4)) Three test suites are built on top of this pattern: - **test_models.py** — full, real-world networks (LeNet, MobileNet, ResNet, …). - **test_operators.py** — one operator at a time (Conv2D, ReLU, BatchNorm, …), parametrized from JSON benchmark configurations. - **test_tensor.py** — round-trip conversion of raw tensors between the two frameworks (dtype, shape, CPU/GPU). A results-analysis script, **analyze_results.py**, post-processes pytest output to classify failures and summarize performance. Components ---------- ModelTester ~~~~~~~~~~~ :class:`ModelTester` (defined in ``utils.py``) is the comparison engine used by every test suite. It owns both models, runs matched forward/backward passes, and asserts that outputs, losses and gradients agree within a tolerance. It can optionally record forward/backward timings and fail the test if Aidge is significantly slower than PyTorch. .. py:class:: ModelTester(model1, model2, name="", test_backward=True, epochs=10, relative_precision=0.001, absolute_precision=0.0001, learning_rate=0.01, cuda=False, benchmark=True, max_slowdown_ratio=10.0, bench_warmup=2, bench_print=True, bench_min_time_ms=0.1) :param torch.nn.Module model1: the reference PyTorch model. :param torch.nn.Module model2: the Aidge model to validate (typically the output of :func:`aidge_interop_torch.wrap`). :param str name: label used in assertion messages and benchmark records. :param bool test_backward: if ``True``, also runs a backward pass, compares the loss, input gradients and parameter gradients, and updates ``model1`` with an SGD step. :param int epochs: number of forward/backward iterations performed by :meth:`test_multiple_step`. :param float relative_precision: relative tolerance used when comparing tensors. :param float absolute_precision: absolute tolerance used when comparing tensors. :param float learning_rate: learning rate used during backward tests. :param bool cuda: run both models on GPU and synchronize before every timing measurement. :param bool benchmark: enable forward/backward timing collection. :param float max_slowdown_ratio: maximum tolerated ``mean(aidge) / mean(torch)`` ratio before :meth:`test_multiple_step` raises. :param int bench_warmup: number of leading epochs excluded from timing (JIT / cache warm-up). :param bool bench_print: print a one-line benchmark summary at the end of :meth:`test_multiple_step`. :param float bench_min_time_ms: skip the slowdown-ratio assertion for a phase whose PyTorch time is below this threshold, since the measurement is overhead-dominated at that scale. .. py:method:: unit_test(input_tensors, label, record_timing=False) Run a single forward pass (and, if ``test_backward``, a backward pass) on both models with the same input tensors, then assert that outputs, loss and gradients match. Raises :class:`AssertionError` on mismatch. .. py:method:: test_multiple_step(input_sizes, label_sizes=None) Run :meth:`unit_test` for ``epochs`` iterations, generating fresh random inputs (and, for the backward pass, all-ones labels) at each step. Wraps any :class:`AssertionError` with the failing epoch number. At the end, prints and/or logs the benchmark summary if ``benchmark=True``. .. py:method:: performance_stats() Return a ``dict`` with ``"forward"`` and ``"backward"`` entries, each containing ``mean1`` (PyTorch), ``mean2`` (Aidge), ``ratio`` and ``samples``. Times are in seconds. .. py:attribute:: BENCH_LOG_FILE :type: str | None :value: "benchmarks.jsonl" Class attribute. Path of a JSONL file that benchmark records are appended to (one JSON object per test: ``name``, ``forward``, ``backward``). Written to disk rather than printed because pytest captures stdout by default. Set to ``None`` to disable file logging. Operator test framework ~~~~~~~~~~~~~~~~~~~~~~~~ ``conftest.py`` provides a small framework for turning an operator plus a JSON benchmark configuration into a fully parametrized pytest fixture. It lets each operator declare, in a few lines, where its test cases come from and how its JSON attributes map onto PyTorch constructor kwargs. .. py:class:: OperatorCase One parametrized test case, as consumed by the fixture built by :func:`make_operator_fixture`. :param str id: the pytest parameter id (shown in ``-k`` filters). :param type op_class: the ``torch.nn.Module`` subclass to instantiate. :param dict torch_kwargs: keyword arguments passed to ``op_class(...)``. :param tuple input_size: shape of the input tensor. :param tuple output_size: shape of the output tensor, or ``None`` to reuse ``input_size``. :param bool test_backward: whether to run and compare the backward pass. :param str expected_failure: if set, the case is registered as ``pytest.mark.xfail`` with this reason instead of failing the suite. .. py:class:: OperatorConfig Full description of one operator under test: either a JSON file to parse, or a hand-written list of :class:`OperatorCase`. :param str name: operator name, used as a prefix for generated test ids. :param type op_class: the ``torch.nn.Module`` subclass under test. :param pathlib.Path json_path: path to a benchmark JSON file (see *JSON configuration formats* below). Ignored if ``manual_cases`` is set. :param list manual_cases: an explicit list of :class:`OperatorCase`, bypassing JSON parsing entirely. :param dict supported_attrs: mapping of JSON attribute name to ``(torch_kwarg_name, type)``, defaults to the module-level ``SUPPORTED_ATTRS`` table (kernel/stride/dilation/padding, ``axes``, ``alpha``, ``beta``, …). :param frozenset ignored_attrs: JSON attribute names to drop for this operator (e.g. an Aidge-only attribute with no PyTorch equivalent). :param dict skip_reasons: mapping of JSON case key to an xfail reason. :param tuple default_input_size: fallback input shape used when a JSON case does not specify one. :param callable specs_to_kwargs: ``(merged_input_specs) -> dict``, used to derive constructor kwargs that depend on tensor shapes (e.g. ``in_channels``/``out_channels`` for ``Conv2D`` or ``FC``). :param str xfail_all: if set, every case for this operator is xfailed with this reason. :param str skip_all: if set, every case is skipped outright (for cases that crash the interpreter, where xfail cannot recover). :param str xfail_backward: if set, only the backward-testing cases are xfailed; forward-only cases still run normally. :param bool forward_only_cases: whether to keep forward-only JSON cases (``shape`` sections); defaults to the module-level ``INCLUDE_FORWARD_ONLY_CASES`` toggle. :param str backend: force ``"cpu"`` or ``"cuda"`` for this operator. .. py:function:: build_cases(cfg) :param OperatorConfig cfg: :returns: the flattened, filtered list of :class:`OperatorCase` for *cfg* (JSON-parsed or manual), with forward-only cases removed unless requested. :rtype: list[OperatorCase] .. py:function:: make_operator_fixture(cfg) Build and return a ``pytest.fixture`` parametrized over every case produced by :func:`build_cases(cfg) `. Each parameter wraps the corresponding PyTorch model with :func:`aidge_interop_torch.wrap`, attaches an Aidge SGD optimizer when the case trains, and applies the right ``xfail``/``skip`` mark (operator-wide skip, operator-wide xfail, backward-only xfail, or per-case xfail, in that priority order). :param OperatorConfig cfg: :returns: a fixture yielding a ``dict`` with keys ``name``, ``torch_model``, ``aidge_model``, ``input_size``, ``output_size``, ``test_backward``. **JSON configuration formats.** ``build_cases`` automatically detects one of the following JSON layouts: ================== ================================================================ Format Structure ================== ================================================================ ``benchmark_ok`` ``base_configuration`` + ``test_configurations.ok.{attr,shape}`` ``benchmark_flat`` ``base_configuration`` + ``test_configurations.{shape,attributes,sweeps,…}`` ``simple`` no ``base_configuration``; flat, named cases ``manual`` no JSON; ``OperatorConfig.manual_cases`` supplies the list ================== ================================================================ For operators whose constructor needs dimensions that live inside the JSON (``Conv2D.in_channels``, ``FC.out_features``, …), supply a ``specs_to_kwargs`` extractor: .. code-block:: python def my_extractor(merged_specs: list[dict]) -> dict: # merged_specs = base input_specs overridden by per-case input_specs return {"in_channels": merged_specs[0]["shape"][1]} MY_OP_CONFIG = OperatorConfig( name="MyOp", json_path=CONFIG_ROOT / "my_op.json", op_class=MyOpModel, specs_to_kwargs=my_extractor, ) Test suites ----------- test_models.py ~~~~~~~~~~~~~~ Runs :class:`ModelTester` on complete networks — ``anomaly_model``, ``kws_model``, ``lenet_model``, ``mobilenet_model``, ``perceptron_model``, ``resnet_model`` — in both eval-only and train configurations. Cases are declared as a plain list of dicts (``test_cases``) and turned into a single parametrized fixture, ``model_test_case``; known-failing configurations carry an ``expected_failure`` string and are marked ``xfail`` automatically. .. code-block:: python pytest test_models.py -v -k "LeNet_eval" test_operators.py ~~~~~~~~~~~~~~~~~~ Runs :class:`ModelTester` operator-by-operator: pooling, activations, ``Conv2D``/``ConvDepthwise2D``/``ConvTranspose2D``, ``FC``, reductions, elementwise ``Add``/``Mul``/``Sub``, ``BatchNorm``, ``LayerNorm``, ``Pad``, and Aidge-specific ops (``ShiftGeLU``, ``ShiftMax``). Each operator is declared once as an :class:`OperatorConfig` — either JSON-backed or with ``manual_cases`` for operators with no benchmark JSON — then turned into a fixture with :func:`make_operator_fixture` and driven by a small ``test__interop`` function. .. code-block:: python pytest test_operators.py -v # all operators pytest test_operators.py -v -k "Conv2D" # one operator pytest test_operators.py -v -k "Conv2D_attributes_kernel_square" pytest test_operators.py -v -k "ReduceMean_axes_1_dim" **Adding a new operator:** 1. Create ``operators/.py`` (a single ``nn.Module``). 2. Export it from ``operators/__init__.py``. 3. Declare an :class:`OperatorConfig` (with ``json_path`` or ``manual_cases``). 4. Call :func:`make_operator_fixture` and write a ``test__interop`` test. test_tensor.py ~~~~~~~~~~~~~~~ Round-trips raw tensors — not models — through :func:`aidge_interop_torch.tensor_converter.torch_tensor_to_aidge` and :func:`aidge_interop_torch.tensor_converter.aidge_tensor_to_torch`, and checks values, dtypes and shapes are preserved. Covers standard dtypes (``float32``/``float64``/``int32``/``int64``), several shapes including empty and scalar tensors, and CPU/GPU transfer (GPU cases are skipped when CUDA or ``aidge_backend_cuda`` is unavailable). .. code-block:: python pytest test_tensor.py -v -k "test_conversion_cpu" Results analysis ----------------- ``analyze_results.py`` parses saved pytest output (verbose or quiet) and the ``benchmarks.jsonl`` log produced by :attr:`ModelTester.BENCH_LOG_FILE`, then prints a categorized report: Aidge bugs to report, performance regressions, test-strategy issues (e.g. shape mismatches), and a per-operator benchmark summary with the ten worst forward/backward slowdown ratios. .. code-block:: console pytest test_operators.py -v 2>&1 | tee results.txt python analyze_results.py results.txt # explicit benchmark file python analyze_results.py results.txt --bench benchmarks.jsonl # or piped directly pytest test_operators.py -v 2>&1 | python analyze_results.py .. py:function:: categorize(rt_msg, ae_msg) :param str rt_msg: the ``RuntimeError`` message captured from a failure, if any. :param str ae_msg: the ``AssertionError`` message captured from a failure, if any. :returns: a ``(category, emoji)`` pair, e.g. ``("outputs_differ", "🔴")`` or ``("shape_mismatch (test strategy)", "🟡")``. :rtype: tuple[str, str] Adding a new interoperability test ----------------------------------- The interoperability framework is designed so that adding support for a new operator or model requires very little boilerplate. Most of the comparison logic is handled by :class:`ModelTester` and the fixtures generated by :func:`make_operator_fixture`. Model tests ~~~~~~~~~~~ Model interoperability tests validate complete neural networks instead of individual operators. To add a new model: 1. **Create the PyTorch model** Add a new model implementation under ``models/``. The model should inherit from ``torch.nn.Module`` and implement the desired network architecture. For example, a LeNet implementation can be added as: .. code-block:: python class TorchLeNet(torch.nn.Module): def __init__(self): ... self.sequence = torch.nn.Sequential(...) def forward(self, x): return self.sequence(x) 2. **Register the model** Add a new entry to the ``test_cases`` list in ``test_models.py``. .. code-block:: python { "name": "LeNet_train", "model": lenet_model.TorchLeNet(), "input_size": (BATCH_SIZE, 1, 32, 32), "output_size": (BATCH_SIZE, 10), "test_backward": True, "expected_failure": "Backward of BatchNorm is not implemented", } Each model configuration defines: - a unique test name; - the instantiated PyTorch model; - the input and output tensor shapes; - whether backward propagation should be validated; - an optional expected failure for known limitations. Once registered, the model is automatically included in the parametrized pytest suite and compared against its Aidge counterpart using :class:`ModelTester`. Operator tests ~~~~~~~~~~~~~~ Adding a new operator generally consists of three steps. 1. **Create a PyTorch wrapper** Create a new file under ``operators/`` containing a minimal ``torch.nn.Module`` that exposes the operator to test. For example, an ``AvgPool2d`` wrapper can be written as: .. code-block:: python class AvgPoolingModel(nn.Module): def __init__( self, kernel_size=(3, 3), stride=(2, 2), padding=(0, 0), ceil_mode=False, count_include_pad=True, ): super().__init__() self.pool = nn.AvgPool2d( kernel_size=kernel_size, stride=stride, padding=padding, ceil_mode=ceil_mode, count_include_pad=count_include_pad, ) def forward(self, x): return self.pool(x) 2. **Declare an** :class:`OperatorConfig` Register the operator in ``test_operators.py`` by creating an :class:`OperatorConfig` describing how the operator should be tested. .. code-block:: python AVGPOOLING_CONFIG = OperatorConfig( name="AvgPooling2D", json_path=CONFIG_ROOT / "avgpooling2d.json", op_class=AvgPoolingModel, ignored_attrs=frozenset( { "rounding_mode", # Aidge-specific "dilations", # not used by AvgPool } ), backend="cuda", ) The configuration specifies: - the operator name used in generated test identifiers; - the JSON benchmark configuration describing the test cases; - the PyTorch wrapper class; - optional ignored or remapped attributes; - backend-specific restrictions, when applicable. 3. **Provide the benchmark configuration** Finally, add the corresponding JSON benchmark file under the benchmark configuration directory. The interoperability framework automatically generates the pytest cases from this configuration using :func:`build_cases` and :func:`make_operator_fixture`.